【GPU使用】如何在物理机和Docker中指定GPU进行推理和训练

张登杰踩2025-03-07 14:23

我的机器上有4张H100卡，我现在只想用某一张卡跑程序，该如何设置。

代码里面设置

bash 复制代码

import os
# 记住要写在impot torch前
os.environ['CUDA_VISIBLE_DEVICES'] = "0, 1"

命令行设置

bash 复制代码

export CUDA_VISIBLE_DEVICES=0,2 # Linux 环境
python test.py # 运行程序

docker 设置

bash 复制代码

docker run --gpus device=2 --shm-size 1024g 

docker run --gpus all --shm-size 1024g 

docker run --gpus device=1,2 --shm-size 1024g

拉取镜像

docker pull vllm/vllm-openai

备注，国内用户，可以使用代理源来加速镜像下载：

docker pull docker.1ms.run/vllm/vllm-openai

无法 import 某个package

bash 复制代码

import sys
sys.path.append('/your/home/package/')

docker设置代理