本文记录pytorch框架中模型的几种状态,主要分为训练和测试两种情况来说。
启用 Batch Normalization 和 Dropout。
不启用 Batch Normalization 和 Dropout。
model.eval()切换到测试模式,在该模式下: 主要用于通知dropout层和BN层在train和validation/test模式间切换: 在train模式下,dropout网络层会按照设定的参数p设置保留激活单元的概率(保留概率=p); BN层会继续计算数据的mean和var等参数并更新。 在eval模式下,dropout层会让所有的激活单元都通过,而BN层会停止计算和更新mean和var,直接使用在训练阶段已经学出的mean和var值。
而with torch.no_grad()则主要是用于停止autograd模块的工作,以起到加速和节省显存的作用。它的作用是将该with语句包裹起来的部分停止梯度的更新,从而节省了GPU算力和显存,但是并不会影响dropout和BN层的行为。 该模式不会影响各层的gradient计算行为,即gradient计算和存储与training模式一样,只是不进行反向传播(back probagation)。