PyTorch训练多任务模型技巧

一、解决在分布式训练中,如果对同一模型进行多次调用的报错

报错如下:

RuntimeError: one of the variables needed for gradient computation has been modified by an inplace operation: [torch.cuda.FloatTensor [256)] is at version 4; expected version 3 instead. Hint: enable anomaly detection to find the operation that failed to compute its gradient, with torch.autograd.set_detect_anomaly(True).

参考知乎文章《【PyTorch踩坑】一个排查了一下午的坑

经过一些调试发现,只有当某些特定情况下才会触发此报错。下面结合一个对比学习的例子(并不是完整的脚本)来简单描述:

python 复制代码
import torch
import torch.nn as nn

from torchvision.models import resnet50

def main():
    model = resnet50(num_classes=256).cuda()
    model = nn.parallel.DistributedDataParallel(model, 
                                                device_ids=[args.local_rank], 
                                                find_unused_parameters=True)
    criterion = nn.MSELoss()
    
    optimizer = torch.optim.SGD(model.parameters(),
                                lr=0.001,
                                momentum=0.99,
                                weight_decay=1e-4)

    for i in range(10):
        input0 = torch.randn((4, 3, 224, 224), dtype=torch.float32).cuda()
        input2 = torch.randn((4, 3, 224, 224), dtype=torch.float32).cuda()

        out1 = model(input0)
        out2 = model(input1)

        loss = criterion(out1, out2)

        optimizer.zero_grad()
        loss.backward()
        optimizer.step()

if __name__ == '__main__':
    main()

经过调试发现,当使用nn.DataParallel并行训练或者单卡训练均可正常运行;另外如果将两次模型调用集成到model中,即通过out1, out2 = model(input0, input1) 的方式在分布式训练下也不会报错。

由此可以猜测:在分布式训练中,如果对同一模型进行多次调用则会触发以上报错,即nn.parallel.DistributedDataParallel方法封装的模型,forword()函数和backward()函数必须交替执行,如果执行多个(次)forward()然后执行一次backward()则会报错。

那么解决此问题的入手点则可以聚焦到nn.parallel.DistributedDataParallel接口上。 通过查询PyTorch官方文档发现此接口下的两个参数:

复制代码
- find_unused_parameters: 如果模型的输出有不需要进行反向传播的,此参数需要设置为True;若你的代码运行后卡住不动,基本上就是该参数的问题。
- broadcast_buffers: 该参数默认为True,设置为True时,在模型执行forward之前,gpu0会把buffer中的参数值全部覆盖到别的gpu上。

问题基本可以定位出来了,即broadcast_buffers=True导致参数被覆盖修改。解决办法:

复制代码
 model = nn.parallel.DistributedDataParallel(model, 
                                             device_ids=[args.local_rank], 
                                             broadcast_buffers=False,
                                             find_unused_parameters=True)

参考

distributed: https://pytorch.org/docs/stable/distributed.html

Inplace error if DistributedDataParallel module that contains a buffer is called twice

相关推荐
我爱一条柴ya9 分钟前
【AI大模型】神经网络反向传播:核心原理与完整实现
人工智能·深度学习·神经网络·ai·ai编程
慕婉030718 分钟前
深度学习概述
人工智能·深度学习
198920 分钟前
【零基础学AI】第30讲:生成对抗网络(GAN)实战 - 手写数字生成
人工智能·python·深度学习·神经网络·机器学习·生成对抗网络·近邻算法
神经星星31 分钟前
新加坡国立大学基于多维度EHR数据实现细粒度患者队列建模,住院时间预测准确率提升16.3%
人工智能·深度学习·机器学习
TY-202531 分钟前
深度学习——神经网络1
人工智能·深度学习·神经网络
cver1232 小时前
CSGO 训练数据集介绍-2,427 张图片 AI 游戏助手 游戏数据分析
人工智能·深度学习·yolo·目标检测·游戏·计算机视觉
FreeBuf_2 小时前
新型BERT勒索软件肆虐:多线程攻击同时针对Windows、Linux及ESXi系统
人工智能·深度学习·bert
强哥之神2 小时前
Meta AI 推出 Multi - SpatialMLLM:借助多模态大语言模型实现多帧空间理解
人工智能·深度学习·计算机视觉·语言模型·自然语言处理·llama
慕婉03074 小时前
Tensor自动微分
人工智能·pytorch·python
神经星星4 小时前
专治AI审稿?论文暗藏好评提示词,谢赛宁呼吁关注AI时代科研伦理的演变
人工智能·深度学习·机器学习