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判断，否则就只输出最后