pytorch中.to(device) 和.cuda()的区别说明

msqyyl

原理.to(device) 可以指定CPU 或者GPU[size=1em]

[size=1em]1 [size=1em]2 [size=1em]3 [size=1em]4 [size=1em]5 [size=1em]6

[size=1em][size=1em]device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu") # 单GPU或者CPU [size=1em]model.to(device) [size=1em]#如果是多GPU [size=1em]if torch.cuda.device_count() > 1: [size=1em]  model = nn.DataParallel(model，device_ids=[0,1,2]) [size=1em]model.to(device)

.cuda() 只能指定GPU[size=1em]

[size=1em]1 [size=1em]2 [size=1em]3 [size=1em]4 [size=1em]5 [size=1em]6 [size=1em]7 [size=1em]8 [size=1em]9

[size=1em][size=1em]#指定某个GPU [size=1em]os.environ['CUDA_VISIBLE_DEVICE']='1' [size=1em]model.cuda() [size=1em]#如果是多GPU [size=1em]os.environment['CUDA_VISIBLE_DEVICES'] = '0,1,2,3' [size=1em]device_ids = [0,1,2,3] [size=1em]net  = torch.nn.Dataparallel(net, device_ids =device_ids) [size=1em]net  = torch.nn.Dataparallel(net) # 默认使用所有的device_ids [size=1em]net = net.cuda()

[size=1em]

[size=1em]1 [size=1em]2 [size=1em]3 [size=1em]4 [size=1em]5 [size=1em]6 [size=1em]7 [size=1em]8 [size=1em]9 [size=1em]10 [size=1em]11 [size=1em]12 [size=1em]13

[size=1em][size=1em]class DataParallel(Module): [size=1em]    def __init__(self, module, device_ids=None, output_device=None, dim=0): [size=1em]        super(DataParallel, self).__init__() [size=1em]        if not torch.cuda.is_available(): [size=1em]            self.module = module [size=1em]            self.device_ids = [] [size=1em]            return [size=1em]        if device_ids is None: [size=1em]            device_ids = list(range(torch.cuda.device_count())) [size=1em]        if output_device is None: [size=1em]            output_device = device_ids[0]

补充：Pytorch使用To方法编写代码在不同设备(CUDA/CPU)上兼容(device-agnostic)

以前版本的PyTorch编写device-agnostic代码非常困难（即，在不修改代码的情况下在CUDA可以使用或者只能使用CPU的设备上运行）。

即设备无关，可以理解为无论什么设备都可以运行您编写的代码。(PS:个人理解，我没有在网上找到专业解释)

PyTorch 0.4.0使代码兼容

PyTorch 0.4.0通过两种方法使代码兼容变得非常容易：

张量的device属性为所有张量提供了torch.device设备。（注意：get_device仅适用于CUDA张量）

to方法Tensors和Modules可用于容易地将对象移动到不同的设备（代替以前的cpu()或cuda()方法）

我们推荐以下模式：

[size=1em]

[size=1em]1 [size=1em]2 [size=1em]3 [size=1em]4 [size=1em]5 [size=1em]6 [size=1em]7

[size=1em][size=1em]# 开始脚本，创建一个张量 [size=1em]device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu") [size=1em]... [size=1em]# 但是无论你获得一个新的Tensor或者Module [size=1em]# 如果他们已经在目标设备上则不会执行复制操作 [size=1em]input = data.to(device) [size=1em]model = MyModule(...).to(device)

以上为个人经验，希望能给大家一个参考，也希望大家多多支持脚本之家。

941287666