python - 如何在 pytorch 中进行卷积，其中每个小批量的内核都不同？

Question

让我用一个例子来表达标题：

令 A 为形状为 [16, 15 , 128, 128] 的张量（表示 [batchsize, channels, height, width]）

令 B 为形状为 [16, 3 , 128, 128] 的张量（表示 [batchsize, channels, height, width]）

我想输出一个形状为 [16, 5 , 128, 128] 的张量（这意味着 [batchsize, channels, height, width]）

其中输出的 5 个通道中的第i_th通道是通过将元素 B 与 A 的 3 个通道的第i_th切片相乘来计算的，并且它们沿通道维度执行求和。

您将如何在 pytorch 中执行该操作？

谢谢！

PD: 很难表达我想从手术中得到什么，如果我不清楚，请问我，我会尝试重新解释

标签： pythonmachine-learningdeep-learningneural-networkpytorch

Answer 1

我认为您正在寻找torch.repeat_interleave帮助您“扩展”B张量以拥有 15 个通道（3 个输入通道中的 5 组）的方法：

extB = torch.repeat_interleave(B, 3, dim=1)  # extend B to have 15 channels
m = A * extB  # element wise multiplication of A with the extended version of B
# using some reshaping and mean we can get the 5 output channels you want
out = m.view(m.shape[0], 5, 3, *m.shape[2:]).mean(dim=2)