如何有效地规范化张量中的元素组



>上下文:

我正在尝试复制Hinton的"带有EM路由的矩阵胶囊"(https://openreview.net/forum?id=HJWLfGWRb(。

在某些时候,会执行卷积运算(从某种意义上说,输出张量连接到输入张量,并且输出张量中的每个元素仅受大小为 K 的 2D 掩码中包含的输入元素的影响(。

形状w_in,w_in的输入张量x哪里

  • w_in=14

映射到形状w_out,w_out,K,K的输入张量x_mapped中间张量哪里

  • 卷积核大小K=3
  • w_out=6,由与stride=2卷积产生

对维度 2 和 3(大小均为K(求和意味着对连接到输出元素的输入元素求和,输出元素的位置由维度 0 和 1 给出。

问题:

如何根据元素组在输入张量x中的位置,有效地将x_mapped中的元素组归一化(为 1(?

例如:
x_mapped(0,0,2,2)
x_mapped(1,0,0,2)
x_mapped(0,1,2,0)
x_mapped(1,1,0,0)
都连接到x(2,2)(公式为i_out*stride + K_index = i_in(。出于这个原因,我希望这 4 个元素的总和为 1。

我想对x_mapped中"连接"到x中同一元素的所有元素组执行此操作。

我可以通过以下方式弄清楚如何做到这一点:

  1. 构建一个字典,其中输入位置作为键,输出元素列表作为值
  2. 在字典上循环,对给定输入位置的列表中的元素求和并将它们除以该总和

但这对我来说似乎真的很低效。

我通过以下方式解决了这个问题:

  1. 创建一个以 2 元组作为键(坐标为x(和x_mapped元素列表作为值的字典。
  2. 对字典进行一次循环,压缩一个字典项的所有元素,然后规范化。

这是代码:

from collections import defaultdict
import torch
ho = 6
wo = 6
stride = 2
K = 3
d = defaultdict(list)
x_mapped = torch.arange(0,ho*wo*K*K).view(ho,wo,K,K).type(dtype = torch.DoubleTensor)
for i_out in range(0,ho):
for j_out in range(0,wo):
for K_i in range(0,K):
for K_j in range(0, K):
i_in = i_out * stride + K_i
j_in = j_out * stride + K_j
d[(i_in, j_in)].append((i_out, j_out, K_i, K_j))
for _ , value in d.items():
ho_list, wo_list, K_i_list, K_j_list = zip(*value)
x_mapped[ho_list, wo_list, K_i_list, K_j_list] = x_mapped[ho_list, wo_list, K_i_list, K_j_list] / torch.sum(
x_mapped[ho_list, wo_list, K_i_list, K_j_list])

最新更新