基于torchvision的数据集CIFAR10，进行卷积神经网络的架构设计，持续更新中

最新推荐文章于 2025-05-23 08:36:32 发布

码代码的芒果

最新推荐文章于 2025-05-23 08:36:32 发布

阅读量216

点赞数 3

文章标签： cnn 深度学习神经网络

本文链接：https://blog.csdn.net/weixin_41750729/article/details/148075000

版权

最近，芒果在研究卷积神经网络，针对于CIFAR10这个数据集，进行实现图像分类

搭建卷积神经网络
编写训练函数
编写预测函数

下面上代码，后面持续更新中…

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

BATCH_SIZE = 8


def create_dataset():
    """
    获取训练集,测试集
    :return:
    """
    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


# 创建图像卷积神经网络类
class ImgCnnClass(nn.Module):
    def __init__(self):
        super().__init__()
        self.conv1 = nn.Conv2d(in_channels=3, out_channels=6, kernel_size=3, stride=1, padding=0)
        self.pool1 = nn.MaxPool2d(kernel_size=2, stride=2)
        self.conv2 = nn.Conv2d(in_channels=6, out_channels=16, 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.output = nn.Linear(84, 10)

    def forward(self, x):
        # 卷积 + relu + 池化
        x = torch.relu(self.conv1(x))
        x = self.pool1(x)
        # 卷积 + relu + 池化
        x = self.pool2(torch.relu(self.conv2(x)))
        # 参1 样本数  参2 列
        x = x.reshape(x.size(0), -1)
        # 全连接层
        x = torch.relu(self.linear1(x))
        x = torch.relu(self.linear2(x))
        x = self.output(x)  # 后续使用CrossEntropyLoss()损失函数，不需要进行softmax操作
        return x


def train(model, train_dataset):
	pass


def evaluate(model, test_dataset):
	pass


if __name__ == '__main__':
    train_dataset, test_dataset = create_dataset()

    img_model = ImgCnnClass()

    train(img_model, train_dataset)

    evaluate(img_model, test_dataset)