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的情况。

感想:

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

你可能感兴趣的:(#,Pytorch)