Pytorch中的矢量张量元素明智乘法

Question

我试图从表示Pytorch中图像的张量中提取亮度，因此我需要将大小为3（针对三个RGB值权重）的向量逐元素乘以表示图像的3xNxN张量，从而最后，我得到了一个N​​xN矩阵，其中张量的三个通道已与向量中给出的权重相加。

我想存在Pytorch操作可以帮助我实现无循环操作，但我还没有找到它们。

Answer 1

您必须重塑3尺寸的RGB向量，才能像这样广播到3xNxN：

rgb = rgb.reshape(-1, 1, 1)

因此它将具有形状(3, 1, 1)

现在，您可以将其与原始图像相乘并沿第一维求和：

result = torch.sum(rgb * image, dim=0)