卷积神经网络实现图像分类

复制代码
# 1.导入依赖包
import torch
import torch.nn as nn
from torchvision.datasets import CIFAR10
from torchvision.transforms import ToTensor
from torchvision.transforms import Compose
import torch.optim as optim
from torch.utils.data import DataLoader
import time
import matplotlib.pyplot as plt
from torchsummary import summary

BATCH_SIZE = 8


# 2. 获取数据集
def create_dataset():
    # 加载数据集:训练集数据和测试数据
    train = CIFAR10(root='data', train=True, transform=Compose([ToTensor()]))
    valid = CIFAR10(root='data', train=False, transform=Compose([ToTensor()]))
    # 返回数据集结果
    return train, valid


# if __name__ == '__main__':
#     # 数据集加载
#     train_dataset, valid_dataset = create_dataset()
#     # 数据集类别
#     print("数据集类别:", train_dataset.class_to_idx)
#     # 数据集中的图像数据
#     print("训练集数据集:", train_dataset.data.shape)
#     print("测试集数据集:", valid_dataset.data.shape)
#     # 图像展示
#     plt.figure(figsize=(2, 2))
#     plt.imshow(train_dataset.data[1])
#     plt.title(train_dataset.targets[1])
#     plt.show()


# 3.模型构建
class ImageClassification(nn.Module):
    # 定义网络结构
    def __init__(self):
        super(ImageClassification, self).__init__()
        # 定义网络层:卷积层+池化层
        self.conv1 = nn.Conv2d(3, 6, stride=1, kernel_size=3)
        self.pool1 = nn.MaxPool2d(kernel_size=2, stride=2)
        self.conv2 = nn.Conv2d(6, 16, stride=1, kernel_size=3)
        self.pool2 = nn.MaxPool2d(kernel_size=2, stride=2)
        # 全连接层
        self.linear1 = nn.Linear(576, 120)
        self.linear2 = nn.Linear(120, 84)
        self.out = nn.Linear(84, 10)

    # 定义前向传播
    def forward(self, x):
        # 卷积+relu+池化
        x = torch.relu(self.conv1(x))
        x = self.pool1(x)
        # 卷积+relu+池化
        x = torch.relu(self.conv2(x))
        x = self.pool2(x)
        # 将特征图做成以为向量的形式:相当于特征向量
        x = x.reshape(x.size(0), -1)
        # 全连接层
        x = torch.relu(self.linear1(x))
        x = torch.relu(self.linear2(x))
        # 返回输出结果
        return self.out(x)


# if __name__ == '__main__':
#     # 模型实例化
#     model = ImageClassification()
#     summary(model, input_size=(3, 32, 32), batch_size=1)

# 4.训练函数编写
def train(model, train_dataset):
    criterion = nn.CrossEntropyLoss()  # 构建损失函数
    optimizer = optim.Adam(model.parameters(), lr=1e-3)  # 构建优化方法
    epoch = 20  # 训练轮数
    for epoch_idx in range(epoch):
        # 构建数据加载器
        dataloader = DataLoader(train_dataset, batch_size=BATCH_SIZE, shuffle=True)
        sam_num = 0  # 样本数量
        total_loss = 0.0  # 损失总和
        start = time.time()  # 开始时间
        # 遍历数据进行网络训练
        for x, y in dataloader:
            output = model(x)
            loss = criterion(output, y)  # 计算损失
            optimizer.zero_grad()  # 梯度清零
            loss.backward()  # 反向传播
            optimizer.step()  # 参数更新
            total_loss += loss.item()  # 统计损失和
            sam_num += 1
        print('epoch:%2s loss:%.5f time:%.2fs' % (epoch_idx + 1, total_loss / sam_num, time.time() - start))
    # 模型保存
    torch.save(model.state_dict(), 'data/image_classification.pth')




def test(valid_dataset):
    # 构建数据加载器
    dataloader = DataLoader(valid_dataset, batch_size=BATCH_SIZE, shuffle=True)
    # 加载模型并加载训练好的权重
    model = ImageClassification()
    model.load_state_dict(torch.load('data/image_classification.pth'))
    model.eval()
    # 计算精度
    total_correct = 0
    total_samples = 0
    # 遍历每个batch的数据,获取预测结果,计算精度
    for x, y in dataloader:
        output = model(x)
        total_correct += (torch.argmax(output, dim=-1) == y).sum()
        total_samples += len(y)
        # 打印精度
    print('Acc: %.2f' % (total_correct / total_samples))


if __name__ == '__main__':
    # 数据集加载
    train_dataset, valid_dataset = create_dataset()
    # 模型实例化
    model = ImageClassification()
    # 模型训练
    # train(model, train_dataset)
    # 模型预测
    test(valid_dataset)
相关推荐
IT猿手12 小时前
基于CNN-LSTM的深度Q网络(Deep Q-Network,DQN)求解移动机器人路径规划,MATLAB代码
网络·cnn·lstm
hyshhhh13 小时前
【算法岗面试题】深度学习中如何防止过拟合?
网络·人工智能·深度学习·神经网络·算法·计算机视觉
Listennnn13 小时前
优雅的理解神经网络中的“分段线性单元”,解剖前向和反向传播
人工智能·深度学习·神经网络
牙牙要健康15 小时前
【目标检测】【深度学习】【Pytorch版本】YOLOV3模型算法详解
pytorch·深度学习·目标检测
MPCTHU16 小时前
预测分析(三):基于机器学习的分类预测
人工智能·机器学习·分类
Scc_hy16 小时前
强化学习_Paper_1988_Learning to predict by the methods of temporal differences
人工智能·深度学习·算法
誉鏐16 小时前
从零开始设计Transformer模型(1/2)——剥离RNN,保留Attention
人工智能·深度学习·transformer
神经星星17 小时前
无需预对齐即可消除批次效应,东京大学团队开发深度学习框架STAIG,揭示肿瘤微环境中的详细基因信息
人工智能·深度学习·机器学习
程序员Linc17 小时前
写给新人的深度学习扫盲贴:向量与矩阵
人工智能·深度学习·矩阵·向量
补三补四17 小时前
机器学习-聚类分析算法
人工智能·深度学习·算法·机器学习