pytorch使用gpu的两种方式

作者：Cpp五条 | 2024-04-05 01:45:00

踩

pytorch使用gpu

在使用gpu进行训练或推理会比纯用cpu快好几倍，所以一般我们如果设备有gpu都尽量会用上gpu。
首先能使用gpu的有：数据（输入的图片、标注的label），损失函数，网络模型。

方法一

这三处都调用.cuda()进行返回。
网络模型：
在这里插入图片描述
损失函数：

数据（输入的图片、标注的label）：
训练集、验证集、测试集（这里拿训练集贴图举例）。

首先定义网络训练的设备，然后三处都调用.to(device)进行返回

# 定义训练的设备
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
1
2

网络模型：
在这里插入图片描述
损失函数：

数据（输入的图片、标注的label）：
训练集、验证集、测试集（这里拿训练集贴图举例）。

声明：本文内容由网友自发贡献，不代表【wpsshop博客】立场，版权归原作者所有，本站不承担相应法律责任。如您发现有侵权的内容，请联系我们。转载请注明出处：https://www.wpsshop.cn/w/Cpp五条/article/detail/362499?site