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

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

.cuda()方法

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)方法 （推荐）

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块显卡，如果要指定显卡，可以参考下面的程序：

# =========== 方法一 ==============
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的内存，按道理迁移时存储地址是会变化的呀）

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()