Logo Questions Linux Laravel Mysql Ubuntu Git Menu
 

Pytorch: can we use nn.Module layers directly in forward() function?

Generally, In the constructor, we declare all the layers we want to use. In the forward function, we define how the model is going to be run, from input to output.

My question is what if calling those predefined/built-in nn.Modules directly in forward() function? Is this Keras function API style legal for Pytorch? If not, why?

Update: TestModel constructed in this way did run successfully, without an alarm. But the training loss will descend slowly compared with the conventional way.

import torch.nn as nn
from cnn import CNN

class TestModel(nn.Module):
    def __init__(self):
        super().__init__()
        self.num_embeddings = 2020
        self.embedding_dim = 51

    def forward(self, input):
        x = nn.Embedding(self.num_embeddings, self.embedding_dim)(input)
        # CNN is a customized class and nn.Module subclassed
        # we will ignore the arguments for its instantiation 
        x = CNN(...)(x)
        x = nn.ReLu()(x)
        x = nn.Dropout(p=0.2)(x)
        return output = x
like image 210
Kuo Avatar asked Jan 08 '20 09:01

Kuo


1 Answers

You need to think of the scope of the trainable parameters.

If you define, say, a conv layer in the forward function of your model, then the scope of this "layer" and its trainable parameters is local to the function and will be discarded after every call to the forward method. You cannot update and train weights that are constantly being discarded after every forward pass.
However, when the conv layer is a member of your model its scope extends beyond the forward method and the trainable parameters persists as long as the model object exists. This way you can update and train the model and its weights.

like image 196
Shai Avatar answered Nov 09 '22 19:11

Shai