pytorch GPU加速

小颖同学要加油哦

于 2021-05-24 10:38:11 发布

阅读量152

点赞数

CC 4.0 BY-SA版权

分类专栏： pytorch 文章标签： pytorch

本文链接：https://blog.youkuaiyun.com/qq_43412810/article/details/117218990

pytorch 专栏收录该内容

3 篇文章

订阅专栏

该博客主要讨论如何将网络模型、损失函数、训练数据和标签有效地迁移到CUDA指定的GPU上进行训练。通过使用`torch.device()`函数，确保在GPU可用时将数据和模型分配到`cuda:0`设备，否则回退到CPU。这一过程对于高效利用GPU资源进行深度学习训练至关重要。

摘要生成于 C知道，由 DeepSeek-R1 满血版支持，前往体验 >

指定GPU的编号
网络，损失函数，训练的数据，标签需要调入GPU中训练
to.(device)

import torch
"""
net:网络
Loss:loss损失函数
data:数据
label:标签
"""
device = torch.device("cuda:0")
print(device)

device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
print(device)

data = data.to(device)