这些张量与梯度相乘的方法

Question

我有一个函数有两个输入：热图和特征图。热图的形状为 (20, 14, 64, 64)，特征图的形状为 (20, 64, 64, 64)。其中 20 是批量大小，14 是关键点的数量。热图和特征图都具有 64x64 的空间维度，并且特征图具有 64 个通道（在第二维上）。

现在我需要将每个热图乘以特征图的每个通道。所以第一个热图必须乘以特征图的所有 64 通道。第二个包含所有频道，依此类推。

之后，我应该有一个形状为 (20, 14, 64, 64, 64) 的张量，我需要在其上应用全局最大池。

现在的问题是我无法创建新的张量来执行此操作，因为必须保留热图和特征图的梯度。

我的实际（缓慢且不保持梯度）代码是：

def get_keypoint_representation(self, heatmaps, features):
    heatmaps = heatmaps[0]
    pool = torch.nn.MaxPool2d(features.shape[2])
    features = features[:, None, :, :, :]
    features = features.expand(-1, 14, -1, -1, -1).clone()

    for i in range(self.cfg.SINGLE_GPU_BATCH_SIZE):
        for j in range(self.cfg.NUM_JOINTS):
            for k in range(features.shape[2]):
                features[i][j][k] = torch.matmul(heatmaps[i][j], features[i][j][k])

    gmp = features.amax(dim=(-1, -2))
    return gmp

任务概览：

Answer 1

给定 hm 形状 (b, k, h, w) 的热图张量和 fm 形状 (b, c, h, w) 的特征张量。

您可以使用单个 einsum 运算符执行此类操作

>>> z = torch.einsum('bkhw,bchw->bkchw', hm, FM)
>>> z.shape
torch.Size([20, 14, 64, 64, 64])

然后使用amax对空间维度进行max-pooling操作：

 >>> gmp = z.amax(dim=(-1,-2))
 >>> gmp.shape
 torch.Size([20, 14, 64])

这些张量与梯度相乘的方法

way for multiplication of these tensors with gradients

python

machine-learning

neural-network

pytorch

tensor