在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初学者的学习笔记
干货满满,欢迎关注~