Python中使用PyTorch指定多块或特定GPU进行训练的方法指南

1. 引入

PyTorch是一个流行的深度学习框架，它提供了方便的API，使得在GPU上进行模型训练变得简单。
然而，有时我们可能需要指定使用多个GPU或者特定的GPU来进行训练，以满足不同的需求。
之所以在python中指定GPU，而不是在shell中指定，是因为在Python中指定GPU可以直接控制GPU的使用，并且可以更精细地控制计算任务，在Python中指定GPU还可以提供更多的灵活性和可扩展性。
本文将介绍如何在Python中使用PyTorch来指定使用特定的GPU或多个GPU进行模型训练。

2. 具体方法

1. 指定使用某一块GPU

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

2. 指定使用某几块GPU（多卡）

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

注意，必须要放在代码开始处指定，比如

正确的写法：环境变量配置写在torch等导入之前

import os
os.environ["CUDA_VISIBLE_DEVICES"] = "2"
from transformers import AutoConfig, AutoModel, AutoTokenizer
import torch

错误的写法：环境变量配置写在torch等导入之后

from transformers import AutoConfig, AutoModel, AutoTokenizer
import torch
import os
os.environ["CUDA_VISIBLE_DEVICES"] = "2"

3. 命令行指定

在命令行中，也可以这样指定：

$ CUDA_VISIBLE_DEVICES=0 python test_qwen_vl.py
$ CUDA_VISIBLE_DEVICES=0,1 python test_qwen_vl.py

4. 参考

https://blog.csdn.net/xys430381_1/article/details/106635977

作者：ybdesire