【PyTorch踩坑】一个排查了一下午的坑

1 年前

Tramac

比较是偷走快乐的贼

问题

RuntimeError: one of the variables needed for gradient computation has been modified by an inplace operation: [torch.cuda.FloatTensor [2048]] 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进行分布式训练时，遇到以上错误（单卡训练以及并行训练时均正常）。

日志的大概意思是用于梯度计算的变量通过 inplace 操作被修改。网上的一些解决方法基本是检查模型定义中是否有 inplace=True 设置以及 += 操作符。但是这两种方案都不能解决遇到的问题。

问题描述

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

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()