利用gpu训练

方式一

在“完整的模型训练”基础上加上如下代码。

'''创建网络模型'''
model = Model()
if torch.cuda.is_available():
    model = model.cuda()

'''损失函数'''
loss_fn = nn.CrossEntropyLoss()
if torch.cuda.is_available():
    loss_fn = loss_fn.cuda()

 '''训练步骤开始'''
    for data in train_dataloader:
        imgs, targets = data
        if torch.cuda.is_available():
            imgs = imgs.cuda()
            targets = targets.cuda()

 '''测试步骤开始'''
    total_test_loss = 0
    total_accuracy = 0
    with torch.no_grad():
        for data in test_dataloader:
            imgs, targets = data
            if torch.cuda.is_available():
                imgs = imgs.cuda()
                targets = targets.cuda()

方式二

在“完整的模型训练”基础上加上如下代码。

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

'''创建网络模型'''
model = Model()
model = model.to(device)

'''损失函数'''
loss_fn = nn.CrossEntropyLoss()
loss_fn = loss_fn.to(device)

 '''训练步骤开始'''
    for data in train_dataloader:
        imgs, targets = data
        imgs = imgs.to(device)
        targets = targets.to(device)

'''测试步骤开始'''
    total_test_loss = 0
    total_accuracy = 0
    with torch.no_grad():
        for data in test_dataloader:
            imgs, targets = data
            imgs = imgs.to(device)
            targets = targets.to(device)