pytorch-使用GPU加速模型训练

1、问题描述

研究手写数字识别时，总是使用CPU训练模型并不是一个很好的选择，GPU的每个核心虽然慢于CPU,但是GPU对图像处理的速度是比CPU要快的多的。

2、解决方案

network = Net()


def train(epoch,train_loader):
 ......
    for batch_idx,(data,target) in enumerate(train_loader):
        output = network(data)
        loss = F.nll_loss(output, target)
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()
  ......
  

network = Net()
network.cuda()
def train(epoch,train_loader):
 ......
    for batch_idx,(data,target) in enumerate(train_loader):
    	#这里要改一下
    	data = Variable(data).cuda()
        target = Variable(target).cuda()
        output = network(data)
        loss = F.nll_loss(output, target)
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()
  ......