当前位置：首页 > article >正文

pytorh学习笔记——cifar10（四）用VGG训练

article 2024/10/26 14:37:44

1、新建train.py，执行脚本训练模型：

import os
import time

import torch
import torch.nn as nn
import torchvision

from vggNet import VGGbase, VGGNet
from load_cifar import train_loader, test_loader
import warnings
import tensorboardX

# 忽略警告
warnings.filterwarnings('ignore')


def main():
    # 定义超参数
    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')  # 使用GPU训练
    # device = 'cpu'  # 使用CPU训练
    print('device:', device)
    batch_size = 128
    learning_rate = 0.1
    num_epoches = 100  # 训练100个epoch

    # 定义模型
    net = VGGNet().to(device)  # 将模型放入GPU

    # 定义损失函数和优化器
    loss_func = nn.CrossEntropyLoss()  # 交叉熵损失函数
    optimizer = torch.optim.Adam(net.parameters(), lr=learning_rate)  # 优化器使用Adam
    scheduler = torch.optim.lr_scheduler.StepLR(optimizer,
                                                step_size=5,
                                                gamma=0.9)  # 学习率衰减, 每5个epoch，学习率乘以0.9

    # 可视化
    log_path = 'logs/vggNet_log'   # 保存日志的文件夹
    if not os.path.exists(log_path):  # 如果log文件不存在，则创建
        os.makedirs(log_path)
    writer = tensorboardX.SummaryWriter(log_path)  # 创建一个writer

    # 训练模型
    # step_n = 0  # 记录训练次数
    for epoch in range(num_epoches):  # 训练num_epoches个epoch
        print('Epoch {}/{}'.format(epoch, num_epoches))
        begin_time = time.time()   # 记录开始时间
        net.train()  # 设置为训练模式
        for idx, (images, labels) in enumerate(train_loader):  # 遍历训练集，共有391个batch，每个batch有128个样本
            images = images.to(device)  # 将图片数据放入GPU
            labels = labels.to(device)  # 将图片标签放入GPU
            outputs = net(images)  # 前向传播
            loss = loss_func(outputs, labels)  # 计算损失
            optimizer.zero_grad()  # 梯度清零
            loss.backward()  # 反向传播
            optimizer.step()  # 优化器更新参数

            # writer.add_scalar('train_loss', loss.item(), global_step=step_n)  # 将loss添加到writer中
            # writer.add_scalar('train correct', 100.0 * correct.item() / batch_size,
            #                   global_step=step_n)  # 将正确率添加到writer中
            # step_n += 1  # 记录训练次数
            if (idx + 1) % 100 == 0:  # 每100个batch打印一次训练信息,每个batch有128个样本，相当于12800个样本打印一次
                _, pred = torch.max(outputs, dim=1)  # 获取预测结果
                correct = pred.eq(labels).cpu().sum()  # 计算正确率
                # pred：神经网络的输出预测张量。
                # labels：通常表示真实的标签。这个张量与 pred 有相同的形状。
                # pred.eq(labels.data)：这个调用会生成一个布尔张量，表示在 pred 中的每个元素是否等于 labels 中的相应元素。结果会是一个同形状的张量，其中的值为 True 或 False。
                # .cpu()：方法用于将张量从 GPU 转移到 CPU。
                # .sum()方法对布尔张量（True 是 1，False 是 0）进行求和，返回 True 值的数量。也就是说，它返回 pred 中与 labels 相等元素的个数。这通常用于计算模型的正确预测数量。
                # 请注意这里的eq()和sum()是torch中的方法，与python自带的eq()、sum()方法略有不同。
                # 详见https://blog.csdn.net/xulibo5828/article/details/143115452
                print('Train Accuracy: {} %'.format(100 * correct / batch_size))
        scheduler.step()  # 更新学习率
        end_time = time.time()  # 记录结束时间
        print('Each train_epoch take time: {} s'.format(end_time - begin_time))  # 打印每个epoch的耗时

        # 测试模型
        sum_loss = 0  # 记录测试损失
        sum_correct = 0  # 记录测试正确率
        net.eval()  # 设置为测试模式
        begin_time = time.time()   # 记录开始时间
        for idx, (images, labels) in enumerate(test_loader):  # 遍历训练集
            images = images.to(device)  # 将图片数据放入GPU
            labels = labels.to(device)  # 将图片标签放入GPU
            outputs = net(images)  # 前向传播
            # loss = loss_func(outputs, labels)  # 计算损失
            _, pred = torch.max(outputs, dim=1)  # 获取预测结果


            if (idx + 1) % 30 == 0:  # 每30个batch打印一次训练信息
                correct = pred.eq(labels).cpu().sum()  # 计算正确率
                # sum_loss += loss.item()  # 测试损失
                sum_correct += correct.item()  # 测试正确率
                print('Test Accuracy: {} %'.format(100 * correct / batch_size))

        # test_loss = sum_loss * 1.0 / len(test_loader)  # 计算测试损失
        # test_correct = sum_correct * 100.0 / len(test_loader.dataset) / batch_size  # 计算测试正确率

        # writer.add_scalar('test_loss', test_loss, global_step=epoch + 1)  # 将loss添加到writer中
        # writer.add_scalar('test correct', test_correct, global_step=epoch + 1)  # 将正确率添加到writer中
        end_time = time.time()  # 记录结束时间
        print('Each test_epoch take time: {} s'.format(end_time - begin_time))   # 打印每个epoch的耗时

    # 保存模型
    torch.save(net.state_dict(), 'vggNet.pkl')  # 保存模型
    print('Finished Training')

    writer.close()   # 关闭writer


if __name__ == '__main__':
    main()

第1个epoch的准确率：

Train Accuracy: 17.96875 %

第20个epoch的准确率：

Train Accuracy: 82.03125 %

第50个epoch的准确率：

Train Accuracy: 87.5 %

没有继续训练。

2、加入可视化的代码：

import os
import time

import torch
import torch.nn as nn
import torchvision

from vggNet import VGGbase, VGGNet
from load_cifar import train_loader, test_loader
import warnings
import tensorboardX

# 忽略警告
warnings.filterwarnings('ignore')


def main():
    # 定义超参数
    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')  # 使用GPU训练
    # device = 'cpu'  # 使用CPU训练
    print('device:', device)
    batch_size = 128
    learning_rate = 0.1
    num_epoches = 100  # 训练100个epoch

    # 定义模型
    net = VGGNet().to(device)  # 将模型放入GPU

    # 定义损失函数和优化器
    loss_func = nn.CrossEntropyLoss()  # 交叉熵损失函数
    optimizer = torch.optim.Adam(net.parameters(), lr=learning_rate)  # 优化器使用Adam
    scheduler = torch.optim.lr_scheduler.StepLR(optimizer,
                                                step_size=5,
                                                gamma=0.9)  # 学习率衰减, 每5个epoch，学习率乘以0.9

    # 可视化
    log_path = 'logs/vggNet_log'   # 保存日志的文件夹
    if not os.path.exists(log_path):  # 如果log文件不存在，则创建
        os.makedirs(log_path)
    writer = tensorboardX.SummaryWriter(log_path)  # 创建一个writer

    # 训练模型
    step_n = 0  # 记录训练次数
    for epoch in range(num_epoches):  # 训练num_epoches个epoch
        print('Epoch {}/{}'.format(epoch, num_epoches))
        begin_time = time.time()   # 记录开始时间
        net.train()  # 设置为训练模式
        for idx, (images, labels) in enumerate(train_loader):  # 遍历训练集，共有391个batch，每个batch有128个样本
            images = images.to(device)  # 将图片数据放入GPU
            labels = labels.to(device)  # 将图片标签放入GPU
            outputs = net(images)  # 前向传播
            loss = loss_func(outputs, labels)  # 计算损失
            optimizer.zero_grad()  # 梯度清零
            loss.backward()  # 反向传播
            optimizer.step()  # 优化器更新参数

            _, pred = torch.max(outputs, dim=1)  # 获取预测结果
            correct = pred.eq(labels).cpu().sum()  # 计算正确率
            # pred：神经网络的输出预测张量。
            # labels：通常表示真实的标签。这个张量与 pred 有相同的形状。
            # pred.eq(labels.data)：这个调用会生成一个布尔张量，表示在 pred 中的每个元素是否等于 labels 中的相应元素。结果会是一个同形状的张量，其中的值为 True 或 False。
            # .cpu()：方法用于将张量从 GPU 转移到 CPU。
            # .sum()方法对布尔张量（True 是 1，False 是 0）进行求和，返回 True 值的数量。也就是说，它返回 pred 中与 labels 相等元素的个数。这通常用于计算模型的正确预测数量。
            # 请注意这里的eq()和sum()是torch中的方法，与python自带的eq()、sum()方法略有不同。
            # 详见https://blog.csdn.net/xulibo5828/article/details/143115452
            writer.add_scalar('train_loss', loss.item(), global_step=step_n)  # 将loss添加到writer中
            writer.add_scalar('train correct', 100.0 * correct.item() / batch_size,
                              global_step=step_n)  # 将正确率添加到writer中
            step_n += 1  # 记录训练次数
            if (idx + 1) % 100 == 0:  # 每100个batch打印一次训练信息,每个batch有128个样本，相当于12800个样本打印一次

                print('Train Accuracy: {} %'.format(100 * correct / batch_size))
        scheduler.step()  # 更新学习率
        end_time = time.time()  # 记录结束时间
        print('Each train_epoch take time: {} s'.format(end_time - begin_time))  # 打印每个epoch的耗时

        # 测试模型
        sum_loss = 0  # 记录测试损失
        sum_correct = 0  # 记录测试正确率
        net.eval()  # 设置为测试模式
        begin_time = time.time()   # 记录开始时间
        for idx, (images, labels) in enumerate(test_loader):  # 遍历训练集
            images = images.to(device)  # 将图片数据放入GPU
            labels = labels.to(device)  # 将图片标签放入GPU
            outputs = net(images)  # 前向传播
            # loss = loss_func(outputs, labels)  # 计算损失
            _, pred = torch.max(outputs, dim=1)  # 获取预测结果


            if (idx + 1) % 30 == 0:  # 每30个batch打印一次训练信息
                correct = pred.eq(labels).cpu().sum()  # 计算正确率
                # sum_loss += loss.item()  # 测试损失
                sum_correct += correct.item()  # 测试正确率
                print('Test Accuracy: {} %'.format(100 * correct / batch_size))

        test_loss = sum_loss * 1.0 / len(test_loader)  # 计算测试损失
        test_correct = sum_correct * 100.0 / len(test_loader.dataset) / batch_size  # 计算测试正确率

        writer.add_scalar('test_loss', test_loss, global_step=epoch + 1)  # 将loss添加到writer中
        writer.add_scalar('test correct', test_correct, global_step=epoch + 1)  # 将正确率添加到writer中
        end_time = time.time()  # 记录结束时间
        print('Each test_epoch take time: {} s'.format(end_time - begin_time))   # 打印每个epoch的耗时

    # 保存模型
    torch.save(net.state_dict(), 'vggNet.pkl')  # 保存模型
    print('Finished Training')

    writer.close()   # 关闭writer


if __name__ == '__main__':
    main()