apex的O0模式与nn.DataParallel的冲突

参数多卡跑的时候出现错误，发现网上没有这个奇妙的错误姿势就上传一下

RuntimeError: Expected all tensors to be on the same device, but found at least two devices, cuda:1 and cuda:0! (when checking argument for argument weight in method wrapper__cudnn_convolution)

字面意义：当跑到其他卡的时候 tensor 在不同的设备上无法完成计算

简化代码时测试发现为apex的参数设置为FP32时会出现错误

model, opt = amp.initialize(model.cuda(), opt, opt_level="O0")

O0：FP32模式 O1：混合加速模式

解决方法：设置为O1模式，opt_level="O1"

简化代码仅作参考

import os
os.environ["CUDA_VISIBLE_DEVICES"] = "3,4,5"
import torch
from torch import nn, optim
import numpy as np
from apex import amp

data = np.random.randint(0, 6, [1600, 5])
dataset = [(np.array(x, dtype='float32'), int(10 < sum(x) < 20)) for x in data]
dataloader = torch.utils.data.DataLoader(dataset, batch_size=16, shuffle=True)


class Md(nn.Module):
    def __init__(self):
        super().__init__()
        self.linear1 = nn.Linear(5, 25)
        self.relu = nn.ReLU()
        self.linear2 = nn.Linear(25, 2)

    def forward(self, inputs):
        x, target = inputs
        x = self.linear1(x)
        x = self.relu(x)
        x = self.linear2(x)
        return x, target


model = Md()
optimizer = optim.SGD(model.parameters(), lr=0.01)
epochs = 100
criterion = nn.CrossEntropyLoss()

model, opt = amp.initialize(model.cuda(), optimizer, opt_level="O1")  # change here
model = torch.nn.DataParallel(model)

for epoch in range(epochs):
    for batched_inputs in dataloader:
        out, target = model(batched_inputs)
        loss = criterion(out, target)
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()
    print("work")
"""
this code will report a runtime error:
RuntimeError: Expected all tensors to be on the same device, but found at least two devices, cuda:0 and cuda:1! 
(when checking argument for argument mat1 in method wrapper_addmm)
it could show that amp can't use with nn.DataParallel
use "O1" ，not "O0" 
"""

实际上用其他多卡方法就行

参考：当代研究生应当掌握的并行训练方法（单机多卡） - 知乎