pytorch使用GPU训练模型

使用GPU训练深度学习模型是经常使用的，GPU训练模型的速度远远超过CPU训练，而多块GPU并行训练的速度远远超过单个GPU训练。
下面学习单块GPU和多块GPU训练模型。

1. 单块GPU

import torch
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
##假设模型是model
model.to(device)  #如果有GPU，将模型转移到GPU上

2. 多块GPU

import torch
import torch.nn as nn
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

if torch.cuda.device_count()>1:
	model = nn.DataParallel(model)
model.to(device)  #将训练分散在多个GPU上

#也可指定哪块GPU
#os.environ['CUDA_VISIBLE_DEVICES'] = '0,1,3'

#或者device_ids = [0,1,3]
#model = nn.DataParalle(model,device_ids=device_ids)

如果上述内容你已经看完了，恭喜你应该已经掌握了使用Pytorch利用GPU训练模型。
下一篇文章将会构建模型，给出具体实例，测试CPU训练模型、单块GPU训练模型、多块GPU训练模型的时间差别。