Pytorch切换GPU或CPU

Pytorch切换GPU或CPU

Pytorch版本：>=0.40

Pytorch默认是在cpu上训练，我们可以通过to(device)方法切换到GPU

to(device)参数

to(device)方法接受的参数为torch.device类型。可以写为

device = torch.device("cuda:0") # 使用gpu0

我习惯这种写法，如果你系统中没有可用的cuda，直接切换到cpu来运行。

device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu") 

数据或模型调用to(device)方法

在训练前切换GPU，数据和模型都是要放到同一块GPU上的，都可以通过调用to(device)来切换显卡。

1.数据对应：torch.Tensor类型

tensor = torch.randn(2, 2)
device = torch.device("cuda:0") # 简写
tensor = tensor.to(device) # 需要赋值
# Outputs:
# tensor([[-0.2136,  1.3254],
#        [ 1.9691, -0.8401]], device='cuda:0')

2.模型对应：torch.module类型

对于线性层，如torch.nn.modules.linear.Linear；对于损失函数，如torch.nn.modules.loss.CrossEntropyLoss。模型中的对象都是属于torch.module这个父类。

fc = nn.Linear(10,10)
device = torch.device("cuda:0") # 简写
fc = fc.to(device)

我们也可以通过get_device()变量，来查看数据(Tensor)所在的设备。好处在于，可以用得到的device给模型设置device。

tensor = torch.randn(2, 2)
device = torch.device("cuda:0") # 简写
tensor = tensor.to(device) # 需要赋值
device1 = tensor.get_device() # 0
fc = nn.Linear(10,10)
fc = fc.to(device1)

假设我们用GPU0进行训练，保存了模型。下一次训练，发现GPU0正在使用，只能用别的GPU，那在读取模型的时候，也得读到对应的GPU.

device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
checkpoint = torch.load('tensor.pt', map_location=device)
model.load_state_dict(checkpoint['model'])

如果直接torch.load，因为原先使用GPU训练，保存的也是GPU的Tensor，它会默认加载到原用的GPU。

参考：

1. [pytorch官网torch.load函数](https://pytorch.org/docs/1.0.0/torch.html?highlight=torch load#torch.load)
2. pytorch中 如何将gpu与gpu、gpu与cpu 在load时相互转化载入，这位大神写得很好，考虑到DataParallel的情况。

感想：

很喜欢李浩源说的话，“我是一只想成为鲲鹏的菜鸟”。我也是个菜鸟，在朝鲲鹏努力！面对新的知识，程序员都应该自己经过思考消化，再输出出来。