How to know the input shape of model

ghost: <p>I looked into <a href="https://github.com/pytorch/examples/blob/26de41904319c7094afc53a3ee809de47112d387/mnist/main.py#L10" rel="noopener nofollow ugc">the mnist example</a>, and print the model (<code>print(model)</code>) which shows</p>
<pre><code class="lang-auto">Net(
  (conv1): Conv2d(1, 32, kernel_size=(3, 3), stride=(1, 1))
  (conv2): Conv2d(32, 64, kernel_size=(3, 3), stride=(1, 1))
  (dropout1): Dropout(p=0.25, inplace=False)
  (dropout2): Dropout(p=0.5, inplace=False)
  (fc1): Linear(in_features=9216, out_features=128, bias=True)
  (fc2): Linear(in_features=128, out_features=10, bias=True)
)
</code></pre>
<p>However, this doesn’t tell us what the input shape is. In mnist, the shape is <code>(1, 1, 28, 28)</code>, but how do we know the input shape from the model definition (let’s say we don’t know the model is for mnist)? I couldn’t find any info about it.</p>
<p>Thanks!</p>

KFrank: <p>Hi Ghost!</p>
<aside class="quote no-group" data-username="ghost" data-post="1" data-topic="208706" data-full="true">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/letter_avatar_proxy/v4/letter/g/f04885/48.png" class="avatar"> ghost:</div>
<blockquote>
<pre><code class="lang-auto">Net(
  (conv1): Conv2d(1, 32, kernel_size=(3, 3), stride=(1, 1))
  (conv2): Conv2d(32, 64, kernel_size=(3, 3), stride=(1, 1))
  (dropout1): Dropout(p=0.25, inplace=False)
  (dropout2): Dropout(p=0.5, inplace=False)
  (fc1): Linear(in_features=9216, out_features=128, bias=True)
  (fc2): Linear(in_features=128, out_features=10, bias=True)
)
</code></pre>
<p>However, this doesn’t tell us what the input shape is.</p>
</blockquote>
</aside>
<p>This only shows us the <code>Module</code>s in <code>model</code> and not the <code>forward()</code><br>
function that “glues” the <code>Module</code>s together and that may well  perform<br>
other non-trivial processing.  Let’s analyze this making the simplest<br>
assumptions about any missing processing.</p>
<p><code>fc1</code> has <code>in_features = 9216</code>.  There must be some sort of <code>flatten()</code><br>
or <code>reshape</code> just before this.  Convolutions work on images of<br>
arbitrary sizes.  The output of <code>conv2</code> has <code>out_channels = 64</code>, so<br>
the output of <code>conv2</code> has shape <code>[batch_size, 64, H, W]</code>.  Assuming<br>
(correctly) that the batch size just goes along for the ride, after the<br>
<code>flatten()</code> operation, we have <code>64 * H * W</code> “features” to pass into<br>
<code>fc1</code>.</p>
<p>Therefore <code>H * W</code> must be <code>9216 / 64 = 144</code>.  Guessing that the<br>
image is square, it would have shape <code>[64, 12, 12]</code>.  (It doesn’t<br>
have to be square; it could have a shape of, say, <code>[64, 9, 16]</code>,)</p>
<p>Each <code>Conv2d</code> layer (no padding, <code>kernel_size = 3</code>, <code>stride = 1</code>)<br>
trims two rows of pixels off of the image.  Therefore the input to<br>
<code>conv1</code> must have shape <code>[batch_size, 1, 16, 16]</code>.</p>
<p>In general, there will be multiple places in a model where the shape<br>
of the tensor is constrained.  (In this case, the input to <code>fc1</code> has to<br>
have <code>in_features = 9216</code>.)  Then you work backwards from the<br>
constraint see what input shapes would be valid for your model.</p>
<aside class="quote no-group">
<blockquote>
<p>In mnist, the shape is <code>(1, 1, 28, 28)</code></p>
</blockquote>
</aside>
<p>Based on our assumptions, <code>[1, 1, 28, 28]</code> wouldn’t be valid for this<br>
model.  After the first two <code>Conv2d</code> layers, that shape would become<br>
<code>[1, 64, 24, 24]</code>.  You could hypothetically have a factor-of-two<br>
downsampling layer between <code>conv2</code> and <code>fc1</code> to take you down to<br>
<code>[1, 64, 12, 12]</code>, but you would typically perform downsampling<br>
between convolutional layers, rather than after them.</p>
<p>(You could also have an <a href="https://pytorch.org/docs/2.4/generated/torch.nn.AdaptiveAvgPool2d.html" rel="noopener nofollow ugc">AdaptiveAvgPool2d (12, 12)</a> between<br>
<code>conv2</code> and <code>fc1</code>, which would be a pretty common way to make the<br>
model much more flexible with respect to input shape.)</p>
<p>Best.</p>
<p>K. Frank</p>

ghost: <p>Thanks for the detailed explanation!</p>