Pytorch实现卷积神经网络(CNN)案例

Title Pytorch实现卷积神经网络(CNN)案例
Framework PyTorch
User wy8817399@vip.qq.com
Id 43
Created 2/8/26, 9:03 PM
Modified 2/9/26, 6:49 AM
Published Yes
Content

深度学习项目的步骤

    1. 准备数据集.

        这里我们用的时候 计算机视觉模块 torchvision自带的 CIFAR10数据集, 包含6W张 (32,32,3)的图片, 5W张训练集, 1W张测试集, 10个分类, 每个分类6K张图片.

        需要单独安装一下 torchvision包, 即: pip install torchvision

    2. 搭建(卷积)神经网络

    3. 模型训练.

    4. 模型测试.

 

卷积层:

    提取图像的局部特征 -> 特征图(Feature Map), 计算方式:  N = (W - F + 2P) // S + 1

    每个卷积核都是1个神经元.

 

池化层:

    降维, 有最大池化 和 平均池化.

    池化只在HW上做调整, 通道上不改变.

 

案例的优化思路:

    1. 增加卷积核的输出通道数(大白话: 卷积核的数量)

    2. 增加全连接层的参数量.

    3. 调整学习率

    4. 调整优化方法(optimizer...)

    5. 调整激活函数...

    6. ...

import torch
import torch.nn as nn
from torchvision.datasets import CIFAR10
from torchvision.transforms import ToTensor
import torch.optim as optim
from torch.utils.data import DataLoader
import time
import matplotlib.pyplot as plt
from torchinfo import  summary
import os

# ===== 超参数集中管理(最佳实践)=====
BATCH_SIZE = 128  # 增大提升训练速度
LR = 1e-4
EPOCHS = 50
WEIGHT_DECAY = 1e-5
NUM_WORKERS = 4  # 利用多核CPU加速数据加载
MODEL_PATH = './model/image_model.pth'


# 1、准备数据集
def create_dataset():
    # 1。获取训练集
    # 参1:数据集路径。 参2:是否训练集,参 3数据预处理 --> 张量数据, 参4?:是否联网下载
    train_dataset = CIFAR10(root='./data', train=True, transform=ToTensor(), download=True)

    test_dataset = CIFAR10(root='./data', train=False, transform=ToTensor(),download=True)

    return train_dataset, test_dataset
# 2、搭建卷积神经网络
class ImageModel(nn.Module):
    # 1、初始化父类成员,搭建神经网络
    def __init__(self):
        # 初始化父类成员
        super().__init__()
        # 第一个卷积层,输入3通道,输出32通道,卷积核大小3*3,步长1,填充1
        self.conv1 = nn.Conv2d(3,32,3,1,1)
        self.bn1 = nn.BatchNorm2d(32)

        # 第一个池化层
        # 窗口大小2*2,步长1,填充0
        self.pool1 = nn.MaxPool2d(2,2,0)

        # 第2个卷积层 输入32通道,输出128通道,卷积核大小3*3,步长1,填充0
        self.conv2 = nn.Conv2d(32,64,3,1,1)
        self.bn2 = nn.BatchNorm2d(64)

        # 第2个池化层
        self.pool2 = nn.MaxPool2d(2,2,0)

        # 第一个人隐藏层(全连接层),自动计算卷积输入数量,输出512,
        self.fc1 = nn.LazyLinear(512)
        self.bn3 = nn.BatchNorm1d(512)
        # 第一个人隐藏层(全连接层),512,输出256
        self.fc2 = nn.Linear(512,256)
        self.bn4= nn.BatchNorm1d(256)
        # 第一个人隐藏层(全连接层),256,输出10
        self.output = nn.Linear(256,10)
        # 添加dropout层
        self.dropout = nn.Dropout(0.5)

    # 定义前向传播
    def forward(self,x):
        # 第一层:卷积层(加权求和) + 激励层(激活函数)+池化层(降维)
        x = self.pool1(torch.relu(self.bn1(self.conv1(x))))

        # 第二层:卷积层(加权求和) + 激励层(激活函数)+池化层(降维)
        x = self.pool2(torch.relu(self.bn2(self.conv2(x))))

        # 第三层: 全连接层(加权求和)+激励层(激活函数)
        # 全连接层只能处理二维数据,所以要将数据进行拉平(8,16,6,6) -> (8,576)
        # 参1:样本数(行数),参2是列数(特征数)。-1表示自动计算
        # 这里的size(0)是x的总长度,8、因为批次是8,不能用x[0],不然会出现形状为 [16, 6, 6] 的张量
        x = x.reshape(x.size(0), -1)
        # print(x.shape)
        x = torch.relu(self.bn3(self.fc1(x)))
        # 添加dropout层
        x = self.dropout(x)
        # 第四层: 全连接层(加权求和)+激励层(激活函数)
        x = torch.relu(self.bn4(self.fc2(x)))
        # 添加dropout层
        x = self.dropout(x)

        # 输出层
        return self.output(x) # 后期用多分类交叉熵函数CrossEntropyLoss = Softmax()激活函数+损失计算

# 3、模型训练
def train(train_dataset):
    # 1、创建数据加载器
    train_dataloader = DataLoader(train_dataset, batch_size=BATCH_SIZE, shuffle=True)
    # 2、创建模型对象
    model = ImageModel().to(device)
    # 3、创建损失函数
    criterion = nn.CrossEntropyLoss()
    # 4、创建优化器
    optimizer = optim.Adam(model.parameters(), lr=LR, weight_decay=WEIGHT_DECAY)

    # ✅ 优化2: 自动创建模型保存目录
    os.makedirs(os.path.dirname(MODEL_PATH), exist_ok=True)

    # 5、开始训练
    for epoch in range(EPOCHS):
        # 总损失
        total_loss = 0.0
        # 总样本数
        total_sample =0.0
        # 预测总正确数
        total_correct = 0
        # 训练开始时间
        start_time = time.time()


        for i, (x, y) in enumerate(train_dataloader):
            # 1、将数据移动到GPU
            x = x.to(device)
            y = y.to(device)
            model.train()
            # 2、前向传播
            y_pred = model(x)
            # 3、计算损失
            loss = criterion(y_pred, y)
            # 4、梯度清零
            optimizer.zero_grad()
            # 5、反向传播
            loss.backward()
            # 6、梯度更新
            optimizer.step()

            #  统计预测正确的个数
            # total_correct += (y_pred.argmax(dim=1) == y).sum().item()
            # print('#' * 30)
            # print(y_pred)
            # print(y)
            # print(torch.argmax(y_pred,dim=-1))
            # print((torch.argmax(y_pred,dim=-1) ==y))
            # print('#' * 30)

            # 预测总正确数
            total_correct += (torch.argmax(y_pred, dim=-1) == y).sum()
            # print(total_correct)

            # 统计总损失
            # loss.item() 返回单个批次的平均损失 乘以批次数,获取总损失
            total_loss += loss.item() * len(y)

            # 统计总样本数
            total_sample += len(y)

        # 打印训练结果
        print(f'轮数:{epoch + 1}, 训练总损失:{total_loss / total_sample}, 训练准确率:{total_correct / total_sample}, 训练时间:{time.time() - start_time}')

    # 保存模型
    torch.save(model.state_dict(), './model/image_model.pth')

# 4、模型测试
def evaluate(test_dataset):
    # 1、创建数据加载器
    test_dataloader = DataLoader(test_dataset, batch_size=BATCH_SIZE, shuffle=False)
    # 2、创建模型对象
    model = ImageModel().to(device)
    # 3、加载模型参数
    model.load_state_dict(torch.load('./model/image_model.pth'))

    # 4、定义变量统计 预测正确的样本数,总样本数
    total_correct = 0
    total_samples = 0

    model.eval()

    # 遍历数据加载器,获取到每批次的数量
    for x, y in test_dataloader:
        # 1、将数据移动到GPU
        x = x.to(device)
        y = y.to(device)
        # 2、前向传播
        y_pred = model(x)
        # 3、统计预测正确的样本数
        total_correct += (torch.argmax(y_pred, dim=-1) == y).sum()

        # 4、统计总样本数
        total_samples += len(y)

    # 打印测试结果
    print(f'测试准确率:{total_correct / total_samples}')

if __name__ == '__main__':
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    print(torch.cuda.get_device_name(0))
    train_dataset, test_dataset = create_dataset()
    # print(train_dataset,train_dataset.data.shape) #(5000,32,32,3)
    # print(test_dataset, test_dataset.data.shape) #(10000,32,32,3)
    #
    # plt.figure(figsize=(2,2))
    # plt.imshow(train_dataset.data[111])
    # plt.show()

    model = ImageModel().to(device)
    # 查看模型参数,参1:模型,参2 (批次,通道,高,宽),参3 设备
    # 卷积层参数计算公式 = 输入通道数*卷积核尺寸*卷积核数量+卷积核数量
    summary(model,input_size=(BATCH_SIZE, 3,32,32),device=device)

    # 模型训练
    train(train_dataset)

    # 模型测试
    evaluate(test_dataset)