What does model.train() do in PyTorch?
Does it call forward()
in nn.Module
? I thought when we call the model, forward
method is being used.
Why do we need to specify train()?
model.train()
tells your model that you are training the model. So effectively layers like dropout, batchnorm etc. which behave different on the train and test procedures know what is going on and hence can behave accordingly.
More details:
It sets the mode to train
(see source code). You can call either model.eval()
or model.train(mode=False)
to tell that you are testing.
It is somewhat intuitive to expect train
function to train model but it does not do that. It just sets the mode.
Here is the code of module.train()
:
def train(self, mode=True):
r"""Sets the module in training mode."""
self.training = mode
for module in self.children():
module.train(mode)
return self
And here is the module.eval
.
def eval(self):
r"""Sets the module in evaluation mode."""
return self.train(False)
Modes train
and eval
are the only two modes we can set the module in, and they are exactly opposite.
That's just a self.training
flag and currently only Dropout
and BatchNorm
care about that flag.
By default, this flag is set to True
.