如何在pytorch中进行卷积，其中每个迷你批的内核不同

Question

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

设A是形状的张量[16，15128128](这意味着[批量大小，通道，高度，宽度](

设B是形状的张量[16，3，128128](这意味着[批量大小，通道，高度，宽度](

我想输出一个形状为[16，5，128128]的张量(意思是[批量大小，通道，高度，宽度](

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

你会如何在pytorch中进行操作？

谢谢！

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

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)