Pytorch中将数据传递到指定设备

引言

在使用Pytorch进行网络训练时，我们通常需要把数据传递到指定的设备（例如GPU）上，关于这方面的博客有很多，但是大家的方法各不相同，对此我进行了一个总结

---

大家方法各不相同的原因是pytorch将以下不同写法解读为指向同一个设备：
GPU：索引 = “cuda:索引” = torch.device
CPU：“cpu” = torch.device
上述torch.device的初始化方法有以下几种，其中也体现着一些等价性

# GPU
torch.device(0)
torch.device("cuda:0")
torch.device("cuda", 0)
# CPU
torch.device("cpu")

总体来说，传递方法共有3种，分别是“device参数”，“to()”和“cuda() & cpu()”

device参数

# gpu
x = torch.ones((3, 3), device=0)
x = torch.ones((3, 3), device="cuda:0")
x = torch.ones((3, 3), device=device(0/"cuda:0"))
# cpu
x = torch.ones((3, 3), device="cpu")
x = torch.ones((3, 3), device=torch.device("cpu"))

to()

# gpu
x = torch.ones((3, 3)).to(0)
x = torch.ones((3, 3)).to("cuda:0")
x = torch.ones((3, 3)).to(torch.device(0/"cuda:0"))
# cpu
x = torch.ones((3, 3)).to("cpu")
x = torch.ones((3, 3)).to(torch.device("cpu"))

cuda() & cpu()

x = torch.ones((3, 3)).cuda()  # 具体索引等于 torch.cuda.current_device()的结果
x = torch.ones((3, 3)).cpu()

指定GPU进行训练

有时候我们的机器有4张GPU，但是我们只想利用其中索引为2，3的卡进行训练怎么办呢？可以在程序开头下入以下环境变量的设置

os.environ["CUDA_VISIBLE_DEVICES"] = "2, 3"

此时，代码中"cuda:0"实际代表现实中索引为2的卡

关于torch.cuda.set_device()

对于这个函数，官方解释是设置当前使用的设备，我以为在开头用该函数设置以后，后面创建的所有数据就会自动放到该设备上，观察下面代码的结果，发现并不是这样。这个函数我感觉挺鸡肋

torch.cuda.set_device(0)
x = torch.FloatTensor([1, 2])
print(torch.cuda.current_device())  # 结果是0
print(x.device)                     # 结果是cpu