Deep-Learning-with-PyTorch

Subclassing nn.Module
207
# In[25]:
model(img.unsqueeze(0))
# Out[25]:
...
RuntimeError: size mismatch, m1:
➥ [64 x 8], m2: [512 x 32] at c:\...\THTensorMath.cpp:940
Admittedly, the error message is a bit obscure, but not too much so. We find references
to linear in the traceback: looking back at the model, we see that only module
that has to have a 512 × 32 tensor is nn.Linear(512, 32), the first linear module after
the last convolution block.
What’s missing there is the reshaping step from an 8-channel 8 × 8 image to a 512-
element, 1D vector (1D if we ignore the batch dimension, that is). This could be
achieved by calling view on the output of the last nn.MaxPool2d, but unfortunately, we
don’t have any explicit visibility of the output of each module when we use
nn.Sequential. 4
8.3 Subclassing nn.Module
At some point in developing neural networks, we will find ourselves in a situation where
we want to compute something that the premade modules do not cover. Here, it is something
very simple like reshaping, 5 ; but in section 8.5.3, we use the same construction to
implement residual connections. So in this section, we learn how to make our own
nn.Module subclasses that we can then use just like the prebuilt ones or nn.Sequential.
When we want to build models that do more complex things than just applying
one layer after another, we need to leave nn.Sequential for something that gives us
added flexibility. PyTorch allows us to use any computation in our model by subclassing
nn.Module.
In order to subclass nn.Module, at a minimum we need to define a forward function
that takes the inputs to the module and returns the output. This is where we define our
module’s computation. The name forward here is reminiscent of a distant past, when
modules needed to define both the forward and backward passes we met in section
5.5.1. With PyTorch, if we use standard torch operations, autograd will take care of the
backward pass automatically; and indeed, an nn.Module never comes with a backward.
Typically, our computation will use other modules—premade like convolutions or
customized. To include these submodules, we typically define them in the constructor
__init__ and assign them to self for use in the forward function. They will, at the
same time, hold their parameters throughout the lifetime of our module. Note that you
need to call super().__init__() before you can do that (or PyTorch will remind you).
4
Not being able to do this kind of operation inside of nn.Sequential was an explicit design choice by the
PyTorch authors and was left that way for a long time; see the linked comments from @soumith at
https://github.com/pytorch/pytorch/issues/2486. Recently, PyTorch gained an nn.Flatten layer.
5
We could have used nn.Flatten starting from PyTorch 1.3.