训练需要修改的基本语句如下:
os.environ['CUDA_VISIBLE_DEVICE'] = "0, 1, 2"
device = torch.device("cuda:0")
device_ids = [0, 1, 2]
net = torch.nn.DataParallel(model, device_ids = device_ids)
net.to(device)
net = net.module
一定要注意不能忘了
net = net.module
否则会报
AttributeError: ‘DataParallel‘ object has no attribute ‘flow‘
类似的错误
但这么写好像会导致在不同卡上分配的资源不均衡,具体解决方法还在探索,学会之后再更新