pytorch 网络train、eval以及在训练与测试时实现不同输出的方法

pytorch 网络train、eval以及在训练与测试时实现不同输出的方法

深度学习网络往往都是通过训练过程将初始参数调优，然后通过训练后的模型进行测试与预测。对于一些场景，比如图像生成、深度图预测、分割等，有可能需要在训练过程中采用多尺度的思想，生成不同尺度的结果并计算loss，但是在测试与预测阶段，实际上只需要输出最后要求的一个尺度（分辨率）下的结果即可。这就需要我们在测试和训练时实现不同的输出。

方法

pytorch的网络模型有一个属性 training，用于判断当前模型是否处于训练阶段，因此，我们在构建模型时，在forward函数的最后判断是否处于训练阶段：

if self.training:
      return out0, out1, out2, out3, out4
    else:
      return out0

如上例所示，输出时判断 self.training是否为真，如果处于训练阶段，就输出多尺度结果，用于多尺度的loss判断，否则就只输出最后所需要的一个结果即可。

在训练阶段，我们需要让模型处于训练：

model.train()

测试或预测时：

model.eval()

实际上模型默认处于训练的状态，因此如果此前没有eval操作，可以不用设置train。

除此之外，训练与测试还有梯度的反向传播过程不同，测试阶段并不需要backward过程，也不需要修改参数，这时可以设置model.eval()或直接使用 torch.no_grad()。
对于BN层及Dropout，训练与测试并不只有梯度的差别，还是需要用model.eval()加以区分。BN与dropout训练测试的区别可以参考这里