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

最新推荐文章于 2024-03-01 17:49:50 发布

UN.Owen.stargazing

最新推荐文章于 2024-03-01 17:49:50 发布

阅读量1k

点赞数

CC 4.0 BY-SA版权

文章标签：深度学习 pytorch 人工智能

本文链接：https://blog.csdn.net/yuriyri/article/details/124889903

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

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=&#