PyTorch 中层输出和标量因子的批量乘法

Question

我正在使用的模型架构具有

k

并行卷积层。这些层的各个输出使用 MLP 层预测权重进行线性组合。该模型将批量图像作为输入并输出批量图像（具有不同的通道数）。

一个简化的

forward()

函数如下：

# the batch input
#   input channels = 3
#   input dimensions = 32x32
# x.size() = [128, 3, 32, 32]
x

# self.mlp = Flatten -> Linear -> LeakyReLU
# conv_weights.size() = [128, 8]
#   batch size = 128
#   8 scalar values for each of the 8 convolution layers
conv_weights = self.mlp(x)

# run convolutions
# self.convs = ModuleList([Conv2d -> BatchNorm2d -> LeakyReLU])
conv_outputs = []
for conv in self.convs:
    # size = [128, 32, 32, 32]
    #   output channels = 32
    #   output dimensions = 32x32
    conv_outputs.append(conv(x))

result = ???

由于批量大小的额外维度，我在尝试将卷积输出与标量因子相乘时遇到了麻烦。没有批量大小，我可以简单地对每个卷积层进行标量乘法。

如何将每个卷积层的输出与 MLP 输出相乘？