【学习笔记】pytorch多gpu

ssh连服务器docker中运行。

可以先查看gpu运行状态,找空的gpu的id,比如是0,2,5,9这四个gpu

watch -n 0.1 nvidia-smi

代码中加入

device = torch.device('cuda')
model = nn.DataParallel(model.to(device))

在运行得时候指定gpu即可(注意DEVICES最后有一个S,被这个坑了)

CUDA_VISIBLE_DEVICES=0,2,5,9  python test.py

 

你可能感兴趣的:(图像分割学习笔记)