原理
.to(device) 可以指定CPU 或者GPU
|
1
2
3
4
5
6
|
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu") # 单GPU或者CPU
model.to(device)
#如果是多GPU
if torch.cuda.device_count() > 1:
model = nn.DataParallel(model，device_ids=[0,1,2])
model.to(device)
|
.cuda() 只能指定GPU
|
1
2
3
4
5
6
7
8
9
|
#指定某个GPU
os.environ['CUDA_VISIBLE_DEVICE']='1'
model.cuda()
#如果是多GPU
os.environment['CUDA_VISIBLE_DEVICES'] = '0,1,2,3'
device_ids = [0,1,2,3]
net = torch.nn.Dataparallel(net, device_ids =device_ids)
net = torch.nn.Dataparallel(net) # 默认使用所有的device_ids
net = net.cuda()
|
|
1
2
3
4
5
6
7
8
9
10
11
12
13
|
class DataParallel(Module):
def __init__(self, module, device_ids=None, output_device=None, dim=0):
super(DataParallel, self).__init__()
if not torch.cuda.is_available():
self.module = module
self.device_ids = []
return
if device_ids is None:
device_ids = list(range(torch.cuda.device_count()))
if output_device is None:
output_device = device_ids[0]
|
补充：Pytorch使用To方法编写代码在不同设备(CUDA/CPU)上兼容(device-agnostic)
以前版本的PyTorch编写device-agnostic代码非常困难（即，在不修改代码的情况下在CUDA可以使用或者只能使用CPU的设备上运行）。
device-agnostic的概念
即设备无关，可以理解为无论什么设备都可以运行您编写的代码。(PS:个人理解，我没有在网上找到专业解释)
PyTorch 0.4.0使代码兼容
PyTorch 0.4.0通过两种方法使代码兼容变得非常容易：
张量的device属性为所有张量提供了torch.device设备。（注意：get_device仅适用于CUDA张量）
to方法Tensors和Modules可用于容易地将对象移动到不同的设备（代替以前的cpu()或cuda()方法）
我们推荐以下模式：
|
1
2
3
4
5
6
7
|
# 开始脚本，创建一个张量
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
...
# 但是无论你获得一个新的Tensor或者Module
# 如果他们已经在目标设备上则不会执行复制操作
input = data.to(device)
model = MyModule(...).to(device)
|
以上为个人经验，希望能给大家一个参考，也希望大家多多支持服务器之家。
原文链接：https://blog.csdn.net/weixin_43402775/article/details/109223794