把数据和模型迁移到GPU: .cuda()方法和.to(device)方法

把数据从CPU迁移到GPU时，可以用.cuda()方法，也可以用.to(device)方法。示例如下。

.cuda()方法

?

1 2 3 4 5 6 7

import torch import time   t = time.time() b = torch.rand([1024, 1024, 10]) b = b.cuda() print('time:{:6.3f}'.format(time.time() - t))  # 输出: time: 0.084

.to(device)方法 （推荐）

?

1 2 3 4 5 6 7

import torch import time   s = time.time() device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu') a = torch.rand([1024, 1024, 10]).to(device) print('time:{:6.3f}'.format(time.time()-s)) # 输出: time: 0.087

.cuda()方法和.to(device)方法耗时基本差不多。

上面的方法默认把数据迁移到第0块显卡，如果要指定显卡，可以参考下面的程序：

?

1 2 3 4 5 6 7 8 9 10 11

# =========== 方法一 ============== gpu = 5 torch.cuda.set_device(gpu) a = torch.rand([1024, 1024, 10]) a = a.cuda()   # =========== 方法二 ============== gpu = 6 device = 'cuda:' + str(gpu) b = torch.randn([1000, 1000, 20]) b = b.to(device)

 

另外，值得注意到的是，张量在GPU和CPU之间的迁移不是in-place操作，而模型在GPU和CPU之间的迁移是in-place操作。示例如下。（我这里有个疑问，in-place操作意味着存储地址没有改变，可是模型在GPU和CPU之间的迁移，存储地址怎么可能不改变？一个是GPU的显存，一个是CPU的内存，按道理迁移时存储地址是会变化的呀）

?

1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17

import torch import torch.nn as nn   device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')   x_cpu = torch.randn(2, 3) print('x_cpu id: ', id(x_cpu))   x_gpu = x_cpu.to(device) print('x_gpu id: ', id(x_gpu))     net_cpu = nn.Linear(32, 10) print('net_cpu id: ', id(net_cpu))   net_gpu = net_cpu.to(device) print('net_gpu id: ', id(net_gpu))

运行结果：

 

参考：

pytorch中数据在CPU与GPU之间的切换

tensor.to() 和 module.to()