函数`)
相关阅读Pytorch基础https://blog.csdn.net/weixin_45791458/category_12457644.html?spm1001.2014.3001.5482torch.cuda.set_device()函数用于设置当前使用的cuda设备在当拥有多个可用的GPU且能被pytorch识别的cuda设备情况下环境变量CUDA_VISIBLE_DEVICES可以影响GPU设备到cuda设备的映射。由于有些情况下可以不显式写出cuda设备的编号此时指的是当前使用的cuda设备默认为cuda0设备如下所示。import torch print(当前cuda设备是, torch.cuda.current_device()) # 获取当前cuda设备 atorch.tensor([1, 2, 3, 4, 5], devicecuda) # 在当前cuda设备上创建一个张量 print(a) devicetorch.device(cuda) # 创建一个当前cuda设备 btorch.tensor([1, 2, 3, 4, 5], devicedevice) # 在当前cuda设备上创建一个张量 print(b) ctorch.tensor([1, 2, 3, 4, 5]).to(device) # 将张量移动到当前cuda设备 print(c) dtorch.tensor([1, 2, 3, 4, 5]).to(cuda) # 将张量移动到当前cuda设备 print(d) etorch.tensor([1, 2, 3, 4, 5]).cuda() # 将张量移动到当前cuda设备 print(e) ftorch.tensor([1, 2, 3, 4, 5]).cuda(device) # 将张量移动到当前cuda设备 print(f) 输出 当前cuda设备是 0 tensor([1, 2, 3, 4, 5], devicecuda:0) tensor([1, 2, 3, 4, 5], devicecuda:0) tensor([1, 2, 3, 4, 5], devicecuda:0) tensor([1, 2, 3, 4, 5], devicecuda:0) tensor([1, 2, 3, 4, 5], devicecuda:0) tensor([1, 2, 3, 4, 5], devicecuda:0)可以看到上面几个例子的在创建和转移张量时指定了没有编号的当前cuda字符串创建和使用的设备也是没有编号的当前cuda设备在使用.cuda()方法时也没有指定编号或者使用了当前cuda设备。使用torch.cuda.set_device()函数可以设置任意一个可用的cuda设备为当前cuda设备这会影响有关当前cuda设备的命令如下所示。import torch torch.cuda.set_device(1) # 设置当前cuda设备编号为1 print(当前cuda设备是, torch.cuda.current_device()) # 获取当前cuda设备 atorch.tensor([1, 2, 3, 4, 5], devicecuda) # 在当前cuda设备上创建一个张量 print(a) devicetorch.device(cuda) # 创建一个当前cuda设备 torch.cuda.set_device(2) # 设置当前cuda设备编号为2 print(当前cuda设备是, torch.cuda.current_device()) # 获取当前cuda设备 btorch.tensor([1, 2, 3, 4, 5], devicedevice) # 在当前cuda设备上创建一个张量 print(b) torch.cuda.set_device(0) # 设置当前cuda设备编号为0 print(当前cuda设备是, torch.cuda.current_device()) # 获取当前cuda设备 ctorch.tensor([1, 2, 3, 4, 5]).to(device) # 将张量移动到当前cuda设备 print(c) torch.cuda.set_device(1) # 设置当前cuda设备编号为1 print(当前cuda设备是, torch.cuda.current_device()) # 获取当前cuda设备 dtorch.tensor([1, 2, 3, 4, 5]).to(cuda) # 将张量移动到当前cuda设备 print(d) torch.cuda.set_device(2) # 设置当前cuda设备编号为2 print(当前cuda设备是, torch.cuda.current_device()) # 获取当前cuda设备 etorch.tensor([1, 2, 3, 4, 5]).cuda() # 将张量移动到当前cuda设备 print(e) torch.cuda.set_device(0) # 设置当前cuda设备编号为0 print(当前cuda设备是, torch.cuda.current_device()) # 获取当前cuda设备 ftorch.tensor([1, 2, 3, 4, 5]).cuda(device) # 将张量移动到当前cuda设备 print(f) 输出 当前cuda设备是 1 tensor([1, 2, 3, 4, 5], devicecuda:1) 当前cuda设备是 2 tensor([1, 2, 3, 4, 5], devicecuda:2) 当前cuda设备是 0 tensor([1, 2, 3, 4, 5], devicecuda:0) 当前cuda设备是 1 tensor([1, 2, 3, 4, 5], devicecuda:1) 当前cuda设备是 2 tensor([1, 2, 3, 4, 5], devicecuda:2) 当前cuda设备是 0 tensor([1, 2, 3, 4, 5], devicecuda:0)从上面的例子中可以看出torch.cuda.set_device()函数可以接受一个正整数参数作为编号改变当前cuda设备。它还可以接受一个device对象作为参数前提是这个device对象是有编号即不能是当前cuda设备。device torch.device(cuda:1) torch.cuda.set_device(device) # 这是可行的 device torch.device(cuda) torch.cuda.set_device(device) # 这是错误的因为device是一个没有编号的当前cuda设备对象 输出 ValueError: Expected a torch.device with a specified index or an integer, but got:cudapytorch官方不建议使用这个函数而是直接显式指明有关cuda设备的编号而不是依赖当前设备如下所示。import torch atorch.tensor([1, 2, 3, 4, 5], devicecuda:1) # 在cuda1设备上创建一个张量 print(a) devicetorch.device(cuda:2) # 创建一个cuda2设备 btorch.tensor([1, 2, 3, 4, 5], devicedevice) # 在cuda2设备上创建一个张量 print(b) ctorch.tensor([1, 2, 3, 4, 5]).to(device) # 将张量移动到cuda2设备 print(c) dtorch.tensor([1, 2, 3, 4, 5]).to(cuda:0) # 将张量移动到cuda0设备 print(d) etorch.tensor([1, 2, 3, 4, 5]).cuda(1) # 将张量移动到cuda1设备 print(e) ftorch.tensor([1, 2, 3, 4, 5]).cuda(device) # 将张量移动到cuda2设备 print(f) 输出 tensor([1, 2, 3, 4, 5], devicecuda:1) tensor([1, 2, 3, 4, 5], devicecuda:2) tensor([1, 2, 3, 4, 5], devicecuda:2) tensor([1, 2, 3, 4, 5], devicecuda:0) tensor([1, 2, 3, 4, 5], devicecuda:1) tensor([1, 2, 3, 4, 5], devicecuda:2)