【pytorch】单机多GPU报错 :Expected all tensors to be on the same device, but found at least two devices
·
单机多GPU训练:
如果不使用太复杂的模型,只需调整一下模型跟数据就可以设置为单机多GPU训练:
import os
from torch.nn.parallel import DataParallel
os.environ['CUDA_VISIBLE_DEVICES'] ='2,3'
model = DataParallel(model) # 调整模型
for step, pack in enumerate(tqdm(train_data_loader)):
# img = pack['img'][0].cuda()
img = pack['img'].data[0].cuda() # 调整数据
gt_sem_seg=pack['gt_semantic_seg'].data[0].cuda()
很多情况下这样就可以了,但是有时候还是会报错:RuntimeError: Expected all tensors to be on the same device, but found at least two devices, cuda:1 and cuda:0
报错的原因及解决思路
只管的报错原因就是两个数据在不通的GPU上,代码中出现了使它们进行运算地方。按理说不会出现这个原因(pytorch给自动分配好了),但有的模型的有些判断语句或者复制语句,不被python多线程调用更新,就会出现问题,比如这里:
for i in self.unlabeled_cats:
assert torch.all(gt_semantic_seg != i), f'Ground-truth leakage! {i}'
for i in self.clip_unlabeled_cats:
assert torch.all(gt_semantic_seg != i), f'Ground-truth leakage! {i}'
程序刚开始运行的时候:unlabeled_cats
、clip_unlabeled_cats
和 gt_semantic_seg
都在cuda0上面,但是在下一个线程调用的时候,gt_semantic_seg
已经是cuda1上面的了,但是unlabeled_cats
和clip_unlabeled_cats
还在是cuda0上。这时候手动更新,将unlabeled_cats
和clip_unlabeled_cats
移动到cuda1上进行判断。比如这里:
for i in self.unlabeled_cats.to(gt_semantic_seg.device):
assert torch.all(gt_semantic_seg != i), f'Ground-truth leakage! {i}'
for i in self.clip_unlabeled_cats.to(gt_semantic_seg.device):
assert torch.all(gt_semantic_seg != i), f'Ground-truth leakage! {i}'
类似的还有其他地方,比如这里也可这样改:
match_matrix = output[unlabeled_idx]
if unlabeled_cats.device==gt_semantic_seg.device: #判断是否移动device设备
gt_semantic_seg[unlabeled_idx] = unlabeled_cats[match_matrix.argmax(dim=1)]
else:
gt_semantic_seg[unlabeled_idx] = unlabeled_cats.to(gt_semantic_seg.device)[match_matrix.argmax(dim=1)]
更多推荐
已为社区贡献5条内容
所有评论(0)