Pytorch训练过程中改变模型参数 requires_grad 属性

如果模型只在一块GPU上跑,该过程非常简单,只需要训练中途迭代model的 parameters,然后改变各个param的requires_grad 属性即可:

for name, param in model.named_parameters():
    logger.info('parameter of %s'%name)
    logger.info('before requires_grad is : %s'%param.requires_grad)
    param.requires_grad = True
    logger.info('after requires_grad is : %s' % param.requires_grad)

但是如果模型是跑在多块GPU上,就要搞清楚pytorch是如何对同一个Module对象进行分布式处理的。

可以肯定的是,pytorch将gpu_0上的模型进行复制,放到其他GPU上,可以使用DataParallel类:

model.to(device)
    if n_gpu > 1:
        model = torch.nn.DataParallel(model)

DataParallel类有一个数据成员:module,可以获得该用于复制的basic module,通过改变该basic module的参数的requires_grad 属性,即可达到模型在平行计算过程中,各个参数参与梯度回传迭代:

for name, param in model.module.named_parameters():
    logger.info('parameter of %s'%name)
    logger.info('before requires_grad is : %s'%param.requires_grad)
    param.requires_grad = True
    logger.info('after requires_grad is : %s' % param.requires_grad)

这里需要确定的是:pytorch是否是只迭代更新gpu_0上的模型的参数,然后将更新后的参数复制到其他gpu上?还是各个gpu单独迭代模型参数?如果是前者,上述操作即可完成训练过程中改变参数requires_grad 属性,如果是后者,那上面的操作只更新了gpu_0上模型的参数属性,其他gpu上的模型参数属性未更新,是错误的。

验证上述猜想:

pytorch在每次前向传播的过程中,都会将主GPU上的模型,分发给各个GPU上,所以,梯度更新只会在主GPU上,更新模型的requires_grad 属性也只需要在主GPU上即可,参考:https://www.cnblogs.com/jfdwd/p/11466126.html

DataParallel并行计算只存在在前向传播

那如何确定哪一块GPU是主GPU?

发布了78 篇原创文章 · 获赞 4 · 访问量 6万+

猜你喜欢

转载自blog.csdn.net/wangxiaosu/article/details/104408130