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

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

代码里面设置

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

命令行设置

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

docker 设置

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

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

docker设置代理