PyTorch中使用指定的GPU_pytorch 多gpu权重是如何决定到哪个设备的

作者：小小林熬夜学编程 | 2024-04-10 10:18:55

踩

pytorch 多gpu权重是如何决定到哪个设备的

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)

声明：本文内容由网友自发贡献，不代表【wpsshop博客】立场，版权归原作者所有，本站不承担相应法律责任。如您发现有侵权的内容，请联系我们。转载请注明出处：https://www.wpsshop.cn/w/小小林熬夜学编程/article/detail/398114