Pytorch 如何利用GPU训练模型

Pytorch 如何利用GPU训练模型

如何配置Pytorch环境，详见我之前的推文：
https://blog.csdn.net/qq_41940277/article/details/125279454

方法一：

在程序中找到以下三类数据，调用“.cuda”即可。
1.网络模型
2.数据（输入、标签）
3.损失函数
示例：

#网络模型
model=Unet()
if torch.cuda().is_available():
	model=model.cudn()

#数据
imgs, targets = data
if torch.cuda().is_available():
	imgs=imgs.cuda()
	targets=targets.cuda()

#损失函数
loss=nn.CrossEntropyLoss()
if torch.cuda().is_available():
	loss=loss.cudn()

方法二：

在程序中找到以下三类数据，定义设备后，调用“.to(device)”即可。
示例：

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

#网络模型
model=Unet()
model.to(device)

#数据
imgs, targets = data
imgs=imgs.to(device)
targets=targets.to(device)

#损失函数
loss=nn.CrossEntropyLoss()
loss.to(device)