When to use model.train() vs model.eval() in Pytoch?

Question

I have a model that is used in a reinforcement learning algorithm for checkers, a la AlphaZero. Similar to that network, mine features batch normalization after each convolution layer. I am aware that this will cause different behavior/output when using .eval() vs .train()

However, I am unsure of when to use eval() vs train(). The model is used at two different points in the algorithm: First, the network is used to generate many games of self-play. Secondly, the network is trained using the positions of theses games, with the evaluation labels taken from the terminal value of the game (-1, 0, +1) and the 'improved policy' labels are taken to be the visit counts after the UCB-tree-search.

It seems to me that when the network is fully trained, I will use .eval(), as that should be 'what the network really thinks'. Therefore, for the games of self-play, I should also use .eval(). Ostensibly, this should result in stronger games of self-play and thus higher quality data. Finally, if I used .eval() in the self-play step I must also use it in the learning phase, otherwise if the network outputs are different the loss won't even be calculated using the actual outputs! I know that the network learns valuable information from .train(), as the batch norm layers learn about the mean/variance of the data universe.

As I type this, I am starting to suspect that I should be using .train() for both the self-play and learning phases. Still that seems wrong as the discrepancy in output between train() and eval() for a given position can be quite large.

Answer 1

I can see arguments to have the self-play phase use both .train() and .eval(), so I had a look at the implementation of facebook's ELF OpenGo and saw that they have the model in eval mode during self-play (see selfplay.py). I would do as they do because their software seems to work.

Answer 2

You can see from the PyTorch documentation that the eval and the train do the same thing. Although they don't explicitly mention it, the documentation is identical:

Sets the module in evaluation mode.

This has any effect only on certain modules. See documentations of particular modules for details of their behaviors in training/evaluation mode, if they are affected, e.g. Dropout, BatchNorm, etc.

---

Additional ideas from this PyTorch forum:

Yes, they are the same. By default all the modules are initialized to train mode (self.training = True). Also be aware that some layers have different behavior during train/and evaluation (like BatchNorm, Dropout) so setting it matters. Also as a rule of thumb for programming in general, try to explicitly state your intent and set model.train() and model.eval() when necessary.

---

Regarding the situations of self-play and then evaluation, I personally would start by doing self-play in .eval() mode in order to have the highest fidelity between self-play and the final execution mode (i.e. a real game). The discrepancy that you speak of should converge as the model does, but I agree it might be problematically large to begin with.

One idea would be to use a warm start... run in .train() mode for a given number of iterations (practice games, or just rounds of self-play), then phase this out as the model hopefully begins to converge.