调整火炬张量通道的大小

Question

我有一个带有3个通道的割炬张量，我希望它是1个通道（所有其他尺寸都应保持不变）。 因此，如果我当前的尺寸是torch.Size([6, 3, 512, 512])，我希望它是torch.Size([6, 1, 512, 512])

我该怎么做？

Answer 1

这能解决您的问题吗？

a = torch.ones(6, 3, 512, 512)
b = a[:, 0:1, :, :]
print(b.size()) # torch.Size([6, 1, 512, 512])