How to set input layer shape in PyTorch/Lightning?

Ask Questions Forum: ask Machine Learning Questions to our readersCategory: PyTorchHow to set input layer shape in PyTorch/Lightning?
Chris Staff asked 6 months ago
1 Answers
Best Answer
Chris Staff answered 6 months ago

It all depends on your data and your layer.
 
Let’s take the MNIST dataset as an example. Having a shape of (28, 28, 1) – i.e. width, height and number of channels – we can pass it to Linear layers as follows:

    self.layers = nn.Sequential(
      nn.Flatten(),
      nn.Linear(28 * 28 * 1, 64)
    )

Here, the Flatten is necessary in order to make the data unidimensional, as Linear layers only support 1D data
 
In the case of a Conv layer, you only have to specify the channels — i.e., 1, in the case of MNIST (here, 5 is the number of feature maps generated by the particular Conv layer; 26 * 26 * 5 is what happens after applying a Conv layer with kernel size 3 and generating these 5 feature maps):

    self.layers = nn.Sequential(
      nn.Conv2d(1, 5, kernel_size=3),
      nn.Flatten(),
      nn.Linear(26 * 26 * 5, 300),
      nn.ReLU(),
      nn.Linear(300, 64),
      nn.ReLU(),
      nn.Linear(64, 10)
    )

Your Answer

18 + 3 =