使用AlexNet进行手写数字识别：项目结构与代码

最新推荐文章于 2024-07-09 21:10:27 发布

原创

最新推荐文章于 2024-07-09 21:10:27 发布 · 2.1k 阅读

31 ·

CC 4.0 BY-SA版权

文章标签：

#深度学习 #卷积神经网络 #python #神经网络

本文介绍如何使用PyTorch复现经典的AlexNet模型，并应用于MNIST手写数字识别任务，包括数据预处理、模型训练及单张图片测试。

一、文件目录结构

abc.ttf是字体文件，用于在图片上标注预测标签

model.py是AlexNet模型代码

train.py用于训练

test.py用于单张图片测试

MNIST文件夹是数据集文件

二、各部分代码

将train.py调整一下结构后，三部分代码内容为：

train.py

'''
复现AlexNet
MNIST 数据集 大小28x28
'''
import torch
import torch.nn as nn
from torchvision.datasets import MNIST
import torchvision.transforms as transform
from torch.utils.data import DataLoader
from model import AlexNet

'''
MNIST数据集下载与准备
'''
data_train = MNIST('./MNIST/data',
    download = True,
    transform = transform.Compose([
        transform.Resize((32,32)),
        transform.ToTensor()
    ])
    )

data_test = MNIST('./MNIST/data',
    train = 'False',
    download = True,
    transform = transform.Compose([