3. Pytorch：cuda

3.1. 使用指定 GPU

  • 直接终端中设定（推荐）

    CUDA_VISIBLE_DEVICES=1 python my_script.py

  • 代码中设定

    import os
os.environ["CUDA_VISIBLE_DEVICES"] = "2"

  • 使用函数 set_device

    import torch
torch.cuda.set_device(1)

3.2. device 切换

对于一个 tensor 对象， cuda() 返回该对象在CUDA内存中的拷贝:

obj = obj.cuda()

对于一个 nn.Module 实例， cuda() 直接将该模型的参数和buffers转移到GPU:

model.cuda()

另外，使用 to(*args, **kwargs) 可以更方便地管理设备。

 1>>> import torch
 2>>> obj = torch.ones((2,5), dtype=torch.float32)
 3>>> device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
 4>>> device
 5device(type='cuda', index=0)
 6>>> obj = obj.to(device, dtype=torch.float32)
 7>>> obj.device
 8device(type='cuda', index=0)
 9
10>>> net = torch.nn.Linear(10,5,bias=True)
11>>> net.to(device)
12>>> net
13Linear(in_features=10, out_features=5, bias=True)
14>>> net.bias.device
15device(type='cuda', index=0)

3.3. 参考资料

  1. PyTorch中使用指定的GPU

  1. pytorch documentation