神经网络——利用GPU训练

有两种方式可以实现

方式1(.cuda)

将网络模型、数据（输入、标注）和损失函数引入cuda()
网络模型：

  if torch.cuda.is_available():
    test = test.cuda()  # 使用GPU

损失函数：

if torch.cuda.is_available():
   loss_fn = loss_fn.cuda()

数据：（训练数据和测试数据）

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

方式2（.to(device)）

首先定义设备：

# 定义训练的设备（cpu：表示在cpu上运行；cuda：表示在gpu上运行）
device = torch.device("cuda")
# 或者使用下面这种形式：
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

将网络模型、数据（输入、标注）和损失函数定义到device上
网络模型：

# 创建网络模型
test = Test()
test.to(device) # 不需要另外赋值

损失函数：

# 创建损失函数
loss_fn = nn.CrossEntropyLoss()
loss_fn.to(device) # 不需要另外赋值

数据：（训练数据集和测试数据集）

imgs, targets = data
imgs = imgs.to(device) # 需要另外赋值
targets = targets.to(device)