【学习笔记】pytorch多gpu

版权声明：本文为博主原创文章，未经博主允许不得转载。 https://blog.csdn.net/zzzpy/article/details/88077168

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