是否可以自动调整 PyTorch 中 torch.nn.Sequential 内 torch.nn.Flatten 之后的层的后续输入大小?

Is it possible to auto-size the subsequent input of a layer following torch.nn.Flatten within torch.nn.Sequential in PyTorch?

如果我有以下模型 class 例如:

class MyTestModel(nn.Module):

    def __init__(self):
        super(MyTestModel, self).__init__()

        self.seq1 = nn.Sequential(
            nn.Conv2d(3, 6, 3),
            nn.MaxPool2d(2, 2),
            nn.Conv2d(6, 16, 3),
            nn.MaxPool2d(2, 2),
            nn.Flatten(),
            nn.Linear(myflattendinput(), 120), # how to automate this?
            nn.ReLU(),
            nn.Linear(120, 84),
            nn.ReLU(),
            nn.Linear(84, 2),
        )
        self.softmax = nn.Softmax(dim=1)

    def forward(self, x):

        x = self.seq1(x)
        x = self.softmax(x)
        return x

我知道,通常您会让数据加载器为模型提供固定大小的输入,因此 nn.Flatten() 之后的图层输入具有固定大小,但是我想知道您是否可以以某种方式自动计算这个?

PyTorch (>=1.8) 有 LazyLinear 推断输入维度。