PyTorch中使用指定的GPU
PyTorch默认使用从0开始的GPU,如果GPU0正在运行程序,需要指定其他GPU。有如下两种方法来指定需要使用的GPU。1. 类似tensorflow指定GPU的方式,使用CUDA_VISIBLE_DEVICES。1.1 直接终端中设定:CUDA_VISIBLE_DEVICES=1 python my_script.py1.2 python代码中设定:import osos.environ["
·
PyTorch默认使用从0开始的GPU,如果GPU0正在运行程序,需要指定其他GPU。
有如下两种方法来指定需要使用的GPU。
1. 类似tensorflow指定GPU的方式,使用CUDA_VISIBLE_DEVICES。
1.1 直接终端中设定:
CUDA_VISIBLE_DEVICES=1 python my_script.py
1.2 python代码中设定:
import os
os.environ["CUDA_VISIBLE_DEVICES"] = "2"
2. 使用函数 set_device
import torch
torch.cuda.set_device(id)
更多推荐
所有评论(0)