pytorch使用gpu的两种方式

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

方法一

这三处都调用.cuda()进行返回。
网络模型:
pytorch使用gpu的两种方式_第1张图片
损失函数:
pytorch使用gpu的两种方式_第2张图片
数据(输入的图片、标注的label):
训练集、验证集、测试集(这里拿训练集贴图举例)。
pytorch使用gpu的两种方式_第3张图片

方式二

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

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

网络模型:
在这里插入图片描述
损失函数:
在这里插入图片描述
数据(输入的图片、标注的label):
训练集、验证集、测试集(这里拿训练集贴图举例)。
pytorch使用gpu的两种方式_第4张图片

你可能感兴趣的:(pytorch,pytorch,深度学习,人工智能)