PyTorch中如何使用指定的GPU

PyTorch默认使用从0开始的GPU，如果GPU0正在运行程序，需要指定其他GPU。

[attach]26[/attach]

有如下两种方法来指定需要使用的GPU。 [list] []使用[b]CUDA_VISIBLE_DEVICES[/b]，类似tensorflow指定GPU的方法。[/] []直接终端中设定：[code]CUDA_VISIBLE_DEVICES=1 python my_script.py[/code][/] []python代码中设定：[code]import os os.environ["CUDA_VISIBLE_DEVICES"] = "2"[/code][/] []使用函数 [b]set_device，[/b]该函数见 pytorch-master\torch\cuda__init__.py。[code]import torch torch.cuda.set_device(id)[/code][/] [/list]   不过官方建议使用[b]CUDA_VISIBLE_DEVICES[/b]，不建议使用 [b]set_device[/b] 函数。