Pytorch利用GPU训练的两种方式

到这里基本上之前学的Pytorch基础知识就顺完了

文章目录

总结干货

# 方式一：注意网络和损失函数无需重新赋值，赋值也没事
    - 网络 
        cifa10_model = Model()
        if torch.cuda.is_available():
            cifa10_model.cuda()

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

    - 输入和标签
        for data in train_dataloader:
        imgs, targets = data
        if torch.cuda.is_available():
            imgs, targets = imgs.cuda(), targets.cuda()
            
    分别调用它们的.cuda()方法
    注意：if torch.cuda.is_available()的使用
    
# 方式二：注意网络和损失函数无需重新赋值，赋值也没事
    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
    device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu') 
    device = torch.device('cuda:1' if torch.cuda.is_available() else 'cpu')
    可以指定哪个显卡，只有一个显卡的时候等同于第一个
    
    - 网络
        cifa10_model = Model()
        cifa10_model.to(device)
    
    - 损失函数
        loss_fn = nn.CrossEntropyLoss()
        loss_fn.to(device)

    - 输入和标签
       for data in train_dataloader:
        imgs, targets = data
        imgs, targets = imgs.to(device), targets.to(device)

方式一

# Coding by ajupyter
# 日期:2022/2/4 11:14
import torch
import torchvision
from torch import nn
from torch.utils.data import DataLoader
from torch.utils.tensorboard import SummaryWriter

from model import Model

# 加载数据
train_data_set = torchvision.datasets.CIFAR10(root='./dataset/',
                                              train=True,
                                              transform=torchvision.transforms.ToTensor(),
                                              download=True)
test_data_set = torchvision.datasets.CIFAR10(root='./dataset/',
                                             train=True,
                                             transform=torchvision.transforms.ToTensor(),
                                             download=True)
# 利用 DataLoader 加载数据
train_dataloader = DataLoader(dataset=train_data_set, batch_size=64)
test_dataloader = DataLoader(dataset=test_data_set, batch_size=64)

# 模型
cifa10_model = Model()
if torch.cuda.is_available():
    cifa10_model.cuda()

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

# 学习率
# learning_rate = 0.01
learning_rate = 1e-2

# 优化器
optimizer = torch.optim.SGD(cifa10_model.parameters(), lr=learning_rate)
writer = SummaryWriter('./all_train_fluent/')

# 定义训练相关参数
epochs = 20
total_train_steps = 0  # 训练总次数
total_test_steps = 0  # 测试总次数

for epoch in range(epochs):
    print(f'-----第{epoch + 1}轮训练-----')
    epoch_total_train_losses = 0.0  # 每一轮训练总损失

    cifa10_model.train()  # 只对dropout等特定层有用
    # 开始训练
    for data in train_dataloader:
        imgs, targets = data
        if torch.cuda.is_available():
            imgs, targets = imgs.cuda(), targets.cuda()

        output = cifa10_model(imgs)
        loss = loss_fn(output, targets)  # tensor(2.3019, grad_fn=<NllLossBackward>)
        # 优化器优化
        # ****1 梯度清零****
        optimizer.zero_grad()
        # ****2 计算梯度****
        loss.backward()
        # ****3 梯度更新****
        optimizer.step()

        total_train_steps += 1
        epoch_total_train_losses += loss.item()

        if total_train_steps % 100 == 0:
            print(f'第{total_train_steps}次训练 损失为：{loss.item()}')
            writer.add_scalar('train_loss', loss.item(), total_train_steps)

    print(f'-----第{epoch + 1}轮训练 训练集总损失{epoch_total_train_losses}-----')

    cifa10_model.eval()  # 只对dropout等特定层有用
    # 开始测试
    epoch_total_test_losses = 0.0  # 每一轮测试总损失
    epoch_total_test_accuracy = 0.0  # 每一轮正确总数
    # 不考虑梯度情况下计算测试集上的损失
    with torch.no_grad():
        for data in test_dataloader:
            imgs, targets = data
            if torch.cuda.is_available():
                imgs, targets = imgs.cuda(), targets.cuda()
            output = cifa10_model(imgs)
            loss = loss_fn(output, targets)
            epoch_total_test_losses += loss.item()

            # 计算正确率
            right_num_tensor = (torch.argmax(input=output, dim=1) == targets).sum()
            epoch_total_test_accuracy += right_num_tensor.item()  # 横向比较

        print(f'-----第{epoch + 1}轮训练 测试集上的损失为{epoch_total_test_losses}-----')
        print(f'-----第{epoch + 1}轮训练 测试集上的正确率为{epoch_total_test_accuracy / len(test_data_set)}-----')

        total_test_steps += 1
        writer.add_scalar('test_loss', epoch_total_test_losses,
                          total_test_steps)
        writer.add_scalar('test_accuract', epoch_total_test_accuracy / len(test_data_set),
                          total_test_steps)
    # 保存模型
    torch.save(cifa10_model.state_dict(), f'cifa10_model-epoch{epoch}-test_loss{epoch_total_test_losses}')

方式二

# Coding by ajupyter
# 日期:2022/2/4 11:14
import torch
import torchvision
from torch import nn
from torch.utils.data import DataLoader
from torch.utils.tensorboard import SummaryWriter

from model import Model

device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
# device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu') 
# device = torch.device('cuda:1' if torch.cuda.is_available() else 'cpu')
'''
RuntimeError: CUDA error: invalid device ordinal 
'''
print(device)
# 加载数据
train_data_set = torchvision.datasets.CIFAR10(root='./dataset/',
                                              train=True,
                                              transform=torchvision.transforms.ToTensor(),
                                              download=True)
test_data_set = torchvision.datasets.CIFAR10(root='./dataset/',
                                             train=True,
                                             transform=torchvision.transforms.ToTensor(),
                                             download=True)
# 利用 DataLoader 加载数据
train_dataloader = DataLoader(dataset=train_data_set, batch_size=64)
test_dataloader = DataLoader(dataset=test_data_set, batch_size=64)

# 模型
cifa10_model = Model()
cifa10_model.to(device)

# 损失函数
loss_fn = nn.CrossEntropyLoss()
loss_fn.to(device)
# 学习率
# learning_rate = 0.01
learning_rate = 1e-2

# 优化器
optimizer = torch.optim.SGD(cifa10_model.parameters(), lr=learning_rate)
writer = SummaryWriter('./all_train_fluent/')

# 定义训练相关参数
epochs = 20
total_train_steps = 0  # 训练总次数
total_test_steps = 0  # 测试总次数

for epoch in range(epochs):
    print(f'-----第{epoch + 1}轮训练-----')
    epoch_total_train_losses = 0.0  # 每一轮训练总损失

    cifa10_model.train()  # 只对dropout等特定层有用
    # 开始训练
    for data in train_dataloader:
        imgs, targets = data
        imgs, targets = imgs.to(device), targets.to(device)
        output = cifa10_model(imgs)
        loss = loss_fn(output, targets)  # tensor(2.3019, grad_fn=<NllLossBackward>)
        # 优化器优化
        # ****1 梯度清零****
        optimizer.zero_grad()
        # ****2 计算梯度****
        loss.backward()
        # ****3 梯度更新****
        optimizer.step()

        total_train_steps += 1
        epoch_total_train_losses += loss.item()

        if total_train_steps % 100 == 0:
            print(f'第{total_train_steps}次训练 损失为：{loss.item()}')
            writer.add_scalar('train_loss', loss.item(), total_train_steps)

    print(f'-----第{epoch + 1}轮训练 训练集总损失{epoch_total_train_losses}-----')

    cifa10_model.eval()  # 只对dropout等特定层有用
    # 开始测试
    epoch_total_test_losses = 0.0  # 每一轮测试总损失
    epoch_total_test_accuracy = 0.0  # 每一轮正确总数
    # 不考虑梯度情况下计算测试集上的损失
    with torch.no_grad():
        for data in test_dataloader:
            imgs, targets = data
            imgs, targets = imgs.to(device), targets.to(device)
            output = cifa10_model(imgs)
            loss = loss_fn(output, targets)
            epoch_total_test_losses += loss.item()

            # 计算正确率
            right_num_tensor = (torch.argmax(input=output, dim=1) == targets).sum()
            epoch_total_test_accuracy += right_num_tensor.item()  # 横向比较

        print(f'-----第{epoch + 1}轮训练 测试集上的损失为{epoch_total_test_losses}-----')
        print(f'-----第{epoch + 1}轮训练 测试集上的正确率为{epoch_total_test_accuracy / len(test_data_set)}-----')

        total_test_steps += 1
        writer.add_scalar('test_loss', epoch_total_test_losses,
                          total_test_steps)
        writer.add_scalar('test_accuract', epoch_total_test_accuracy / len(test_data_set),
                          total_test_steps)
    # 保存模型
    torch.save(cifa10_model.state_dict(), f'cifa10_model-epoch{epoch}-test_loss{epoch_total_test_losses}')