If I have the following model class for example:
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
I know, normally you would let the data loader give a fixed size input to the model, thus having a fixed size for the input of the layer after
nn.Flatten(), however I was wondering if you could somehow compute this automatically?
PyTorch (>=1.8) has LazyLinear which infers the input dimension.