Pytorch使用GPU加速

最新推荐文章于 2024-07-23 14:52:41 发布

海马走角

最新推荐文章于 2024-07-23 14:52:41 发布

阅读量920

点赞数

分类专栏：笔记文章标签：深度学习人工智能 pytorch gpu

本文链接：https://blog.csdn.net/qq_39187959/article/details/118726920

版权

笔记专栏收录该内容

12 篇文章 2 订阅

订阅专栏

Pytorch使用GPU加速

'''
在进行深度学习开发时，GPU加速可以提升我们的开发效率。

在PyTorch中使用GPU和TensorFlow中不同，
在TensorFlow如果不对设备进行指定时，TensorFlow检测到GPU就会把自动将数据与运算转移到GPU中。
而PyTorch需要显性的指定数据和运算放在哪里执行，这样的操作比较自由，却也有些繁琐。

'''

import torch

# 检测是否可以使用GPU，使用一个全局变量use_gpu，便于后面操作使用
use_gpu = torch.cuda.is_available()

# 构建网络时，把网络，与损失函数转换到GPU上
model = get_model()
loss_f = torch.nn.CrossEntropyLoss()
if (use_gpu):
    model = model.cuda()
    loss_f = loss_f.cuda()

# 训练网络时，把数据转换到GPU上
if (use_gpu):
    x, y = x.cuda(), y.cuda()

# 取出数据时，需要从GPU转换到CPU上进行操作
if(use_gpu):
    loss = loss.cpu()
    acc = acc.cpu()