n*m*3 输入图像到 PyTorch 中的 nxm 标签

Question

我的网络输入是一个尺寸为 nm 的 RGB 图像，我怎样才能使输出尺寸为 nm。

class Net(nn.Module):
    def __init__(self):
        super(Net, self).__init__()
        self.conv1 = nn.Conv2d(3, 20, kernel_size = 5)
        self.conv2 = nn.Conv2d(20, 50, kernel_size = 3)
        self.conv3 = nn.ConvTranspose2d(50,20, kernel_size = 5)
        self.conv4 = nn.ConvTranspose2d(20,1, kernel_size = 3)

    def forward(self, x):
        x = F.relu(self.conv1(x))
        x = F.relu(self.conv2(x))
        x = F.relu(self.conv3(x))
        x = F.relu(self.conv4(x))
    return x

我目前输出一个1 * n * m。如何输出 n*m？

Answer 1

如果你想将一个张量reshape成不同大小但元素数量相同，一般可以使用torch.view.

对于您的情况，有一个更简单的解决方案：torch.squeeze returns 删除所有尺寸为 1 的张量。