pytorch-使用GPU加速模型训练

在安全路上的白白 lv.1

发布时间：2023-05-04 20:52:50 299

文章目录

1、问题描述
2、解决方案

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

文章来源： https://blog.51cto.com/u_14597003/5988983

特别声明：以上内容（图片及文字）均为互联网收集或者用户上传发布，本站仅提供信息存储服务！如有侵权或有涉及法律问题请联系我们。