在 nn.sequential 中編寫張量視圖層

Question

我有一個sequential容器，我想在里面使用Tensor.view函數。 因此，我當前的解決方案如下所示：

class Reshape(nn.Module):
    def __init__(self, *args):
        super().__init__()
        self.my_shape = args

    def forward(self, x):
        return x.view(self.my_shape)

在我的AutoEncoder類中，我有：

self.decoder = nn.Sequential(
                torch.nn.Linear(self.bottleneck_size, 4096*2),
                Reshape(-1, 128, 8, 8),
                
                nn.UpsamplingNearest2d(scale_factor=2), 
                ...

有沒有辦法直接在sequential塊中重塑張量，這樣我就不需要使用外部創建的Reshape類？ 謝謝

Answer 1

您可以使用 Pytorch 文檔中的UNFLATTEN層：

Unflattens tensor dim 將其擴展為所需的形狀。 與順序一起使用。

所以你會有：

self.decoder = nn.Sequential(
            torch.nn.Linear(self.bottleneck_size, 4096*2),
            nn.Unflatten(1, (1, 128, 8, 8)), # The first parameters is the dimension you would like to unflatten, note that dimension 0 is usually your batch size. So here we need dimension 1.
            # These alsos work
            # nn.Unflatten(1, (-1, 128, 8, 8)), 
            # nn.Unflatten(1, (128, 8, 8)),                 nn.UpsamplingNearest2d(scale_factor=2), 
            ...

如果您還沒有，還應該在Pytorch 論壇上查看此討論。 這里還有 torchvision 模型過去是如何在 Pytorch 中實現的。 你可以看到他們已經將Tensor.view從其余的Sequential模塊中分離出來，並將其應用在forward中。 相同代碼的當前版本現在使用flatten ，這意味着在這里使用unflatten是合理的。