Should model be in train mode when saving it?

vision
1

When saving a model like the following snip of code, does it then matter if the model is in train() or eval() mode? Does train() and eval() change the model.state_dict()? And if so, should model.train() be called before saving the model?

for epoch in range(5):
    # prepare model for training
    model.train()
    for i, (local_batch, local_labels) in enumerate(training_generator):
        ....

    # prepare model for evaluation
    model.eval()
    for i, (local_batch, local_labels) in enumerate(validation_generator):
        ....
    
    torch.save(model.state_dict(),  os.path.join('models', f'model_{epoch+1:03}.pt'))
2

No, it doesn’t. Calling model.train() or model.eval() will change the internal self.training flag to True or False, respectively. This flag is then used in some modules (e.g. nn.Dropout and batchnorm layers) so switch the behavior in their forward method.

1 Like
3

Do you mean that I should put train() in this function?

def forward(self, x):
    return self.model(x)
4

No, you would call model.train() and model.eval() before starting the training and evaluation once.