使用指定编号的GPU进行模型的训练和推理,并且用gpustat进行监控显存

在train.py 或者infer.py的代码文件输入以下代码即可
图中代码所指定的所指定的GPU为GPU：0

import os
os.environ['CUDA_VISIBLE_DEVICES'] = '0'

且在GPU的训练语句中将“use_gpu”调整为：

use_gpu = True

在Linux终端输入

pip install gpustat

安装后在终端输入：

gpu-stat -cpu

就可以实时查看Linux下GPU显存的使用情况了。

博主微信公众号：
AI初学者的学习笔记
干货满满，欢迎关注~