What does model.train() do in PyTorch?

python pytorch
Aerin picture Aerin · Jul 20, 2018 · Viewed 52k times · Source

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()?

Answer

Umang Gupta picture Umang Gupta · Jul 20, 2018

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.

Related questions

How to check if pytorch is using the GPU?

I would like to know if pytorch is using my GPU. It's possible to detect with nvidia-smi if there is any activity from the GPU during the process, but I want something written in a python script. Is there a …

Read More
How to initialize weights in PyTorch?

How to initialize the weights and biases (for example, with He or Xavier initialization) in a network in PyTorch?

Read More
How does the "view" method work in PyTorch?

I am confused about the method view() in the following code snippet. class Net(nn.Module): def __init__(self): super(Net, self).__init__() self.conv1 = nn.Conv2d(3, 6, 5) self.pool = nn.MaxPool2d(2,2) self.conv2 = nn.Conv2d(6, 16, 5) self.fc1 = …

Read More