Search code examples
pythondeep-learningneural-networkpytorchconv-neural-network

Pytorch: Automatically determine the input shape of Linear layer after Conv1d


I want to build a model with several Conv1d layers followed by several Linear layers. Conv1d layers will work for data of any given length, the problem comes at the first Linear layer, because the data length is unknown at initialization time. Every time the length of the input data changes, the output size of Conv1d layers will change, hence a change in the required in_features of the first Linear layer.

Note: I learned CNN and I am aware of how to calculate the output dimensions by hand. I am looking for a programmatic way to determine it, because I have to experiment many times with different length of input data.

Question: In pytorch, how do you automatically figure out the output dimension after many Conv1d layers and set the in_features for the following Linear layer?


Solution

  • You can use the builtin nn.LazyLinear which will find the in_features on the first inference and initialize the appropriate number of weights accordingly:

    linear = nn.LazyLinear(out_features)