Initialize module only upon first input? (coding best practices)

I see the issues. A cleaner way would be to set the in_features as the shape calculations, e.g. in_features=512*7*7. So that it would be clear the last conv layer outputs 512 channels and a volume of spatial size 7x7.

This can be easily done by passing a random input to the conv layer:

nn.Conv2d(3, 6, 3, 1, 1)(torch.randn(1, 3, 24, 24)).shape

However, you would still need to know the input shape of your image.
I’m not sure there is a clean way to achieve the shape inference currently.
Maybe someone else has a good idea.

1 Like