请选择 进入手机版 | 继续访问电脑版
查看: 1448|回复: 0

PyTorch中使用指定的GPU

[复制链接]

665

主题

1234

帖子

6557

积分

xdtech

Rank: 5Rank: 5

积分
6557
发表于 2019-6-11 18:29:24 | 显示全部楼层 |阅读模式
本帖最后由 shaoheshaohe 于 2019-6-11 18:30 编辑

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 osos.environ["CUDA_VISIBLE_DEVICES"] = "2"

2. 使用函数 set_device

import torchtorch.cuda.set_device(id)

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

官方建议使用CUDA_VISIBLE_DEVICES,不建议使用 set_device 函数。



回复

使用道具 举报

您需要登录后才可以回帖 登录 | 立即注册

本版积分规则

快速回复 返回顶部 返回列表