欢迎您访问程序员文章站本站旨在为大家提供分享程序员计算机编程知识!
您现在的位置是: 首页

PyTorch指定GPU

程序员文章站 2022-06-17 07:59:00
...

PyTorch默认使用从0开始的GPU,如果想用其他GPU,需要指定。

有两种方法可以指定需要用的GPU:

1.类似tensorflow指定GPU的方式,使用CUDA_VISIBLE_DEVICES
1.1 直接终端中设定:
CUDA_VISIBLE_DEVICES=1 python main.py
1.2 python代码中设定:
import os
os.environ["CUDA_VISIBLE_DEVICES"] = "1"
2.使用函数 set_device
import torch
torch.cuda.set_device(id)

该函数见 pytorch-master\torch\cuda__init__.py。

不过官方建议使用CUDA_VISIBLE_DEVICES,不建议使用 set_device 函数。
其他详情可参见官方文档:https://pytorch.org/docs/master/cuda.html
https://pytorch-cn.readthedocs.io/zh/stable/package_references/torch-cuda/(中文版)

相关标签: GPU