weixin_42419251的博客

原创 Pytorch指定GPU进行训练的方法

Pytorch指定GPU进行训练的方法   在用Pytorch深度学习框架进行训练模型的时候，默认使用服务器的第“0”号GPU训练。当“0”号GPU被占用时，不能正常进行训练，通常会报错“CUDA out of memory”，这时需要选用其余GPU进行训练。用os模块即可解决：   如上图所示，此时第“0”号GPU已被占用，则需要在用于训练的.py文件中加入如下的代码(一般在import之后加入)： import os #"1, 2"表示训练的时候选用两块GPU，优先选用"1"号GPU os.envi