How to customize model.train() and model.eval() on Pytorch

1

model.train() and model.eval() are flags that tell the model that you are training the model and testing mode respectively. This will make the model behave accordingly to techniques such us dropout that have different procedures in train and testing mode.

class NeuralNet(torch.nn.Module):
    def __init__(self, input_size, hidden_size, num_classes):
        super(NeuralNet, self).__init__()
        # Code..

    def forward(self, x):
        out = some_func(x)
         ## Code..
        return out

Suppose I’ve written a different function some_func() that I want to be executed only on training mode. How can I modify the functions model.train() and model.eval() to do that ?

1 Like
2

model.train() and model.eval() sets the calling nn.Module's and its children’s modules training attribute to True and False respectively.
If you need to invoke functions based on training or testing mode, you can simply use the network’s training attribute.

def forward(self, x):
    if self.training:
        out = self.train_module(x)
    else:
        out = self.eval_module(x)
    return out
2 Likes
3

Mmm, interesting. I tried that, but it’s not clear for me how to run the code accordingly. Thank you for your answer.