在服务器上指定GPU编号进行训练模型(终端指令操作与代码文件内编写均有)

在pytorch的框架安装后,使用服务器进行模型训练,一般的服务器都是多卡的训练,一般来说pytorch框架时默认GPU:0 进行模型训练,如果在一个已经满了的显存上继续进行模型训练,会报错“out of memory”.因此需要进行换卡执行操作。有以下几种办法:
Linux终端中执行

CUDA_VISIBLE_DEVICES=1 python train.py

以上表示在GPU:1上执行train.py
以下的指令功能相同,指示执行的方法不同;
在train.py添加以下代码:

import os
os.environ["CUDA_VISIBLE_DEVICES"] = "1"

只需要根据实际的占用情况,指定GPU的编号即可,从GPU:0开始这样可以共用服务器了。

猜你喜欢

转载自blog.csdn.net/baidu_39629638/article/details/107770840
今日推荐