可学习的张量加权和

Question

我正在尝试使用PyTorch在编码器-解码器体系结构中实现深度监控策略。其思想是对三个卷积层的结果进行加权和（使用可学习的参数Wi）

假设我们有三个形状相同的张量：A、B和C：(64, 48, 48, 48)

我的目标是对这三个张量做一个加权线性和：(w0 * A + w1 * B + w2 * C)和w0，w1，w2应该是网络可以学习的参数

也许我必须使用torch.nn.Linear(in_features, out_features)，但我不知道在这种情况下会有什么输入和输出特性

有什么建议吗

Answer 1

您可以定义一个自定义参数张量并将w_i存储在其中。然后用这些权重计算矩阵的加权和

按如下方式注册自定义参数：

W = nn.Parameter(torch.rand(3))

您可以手动计算总和：

w0, w1, w2 = W
res = w0*A + w1*B + w2*B 

或者使用^{}来简洁。请注意，此方法不取决于线性和中的组件数量：

X = torch.stack([A, B, C])
res = torch.einsum('mbchw,m->bcw', X, W)

---

你指出的好东西^{}，你可以用这个图层把它画下来。请注意nn.Linear可以将n维张量作为输入：(batch_size, *, in_features)并将输出(batch_size ,*, out_features)，其中*可以是任意数量的维度。在您的例子中in_features是权重的数量out_features是1。换个角度看：你只需要1神经元来计算加权和。但重要的一点是，“特征”维度必须是最后一个维度，即堆栈必须在最后一个维度上完成：

W = nn.Linear(3, 1)
res = W(torch.stack([A, B, C], dim=-1))

输出形状为：

>>> res.shape
torch.Size([64, 48, 48, 48, 1])