正向矢量化Pytork操作

Question

我编写这个PyTorch模块是为了获取两组由两个单独的卷积层输出的特征图，并对它们进行平均。输入是两个张量，包含相同数量的图像和相同数量的通道（[批次大小、通道、高度、宽度]）。两幅图像的特征图大小不同

    class Avg_Combine(nn.Module):
        def __init__(self):
            super(Avg_Combine, self).__init__()
            
        def forward(self, a, b):
            if a.shape[2] < b.shape[2]:
                factor = b.shape[2] / a.shape[2]
                a = functional.interpolate(a, scale_factor=factor)
            elif a.shape[2] > b.shape[2]:
                factor = a.shape[2] / b.shape[2]
                b = functional.interpolate(b, scale_factor=factor)
            
            idx0 = 0
            for img in a:
                idx1 = 0
                for channel_a in img:
                    channel_a = channel_a.add(b[idx0][idx1]) / 2
                    idx1 += 1
                idx0 += 1
        
            return a  

例如，a=[20128,8,8]，b=[20128,4,4]。代码将向上采样b以匹配a的大小。既然a和b的大小都是[20,128,8,8]，那么它们将在“通道”维度上进行平均，以生成大小为[20,128,8,8]的输出

然而，在训练过程中，这个功能比相同层次的普通CNN慢50~100倍。我怀疑for循环是问题的原因。如何将此操作矢量化？关于如何加快这件事，还有其他建议吗

非常感谢！任何建议都将不胜感激