PyTorch搭建LeNet-5模型（在MNIST数据集上准确率接近100%）

最新推荐文章于 2024-11-18 21:59:35 发布

原创

最新推荐文章于 2024-11-18 21:59:35 发布 · 3.3k 阅读

15 ·

CC 4.0 BY-SA版权

文章标签：

#pytorch #python #深度学习

该博客介绍了如何在PyTorch中构建经典的LeNet-5模型，应用于MNIST数据集进行手写数字识别。步骤包括数据预处理、模型定义、训练过程以及测试评估，最终模型达到了99.09%的测试准确率。

PyTorch搭建LeNet-5模型

https://www.cnblogs.com/gshang/p/13099170.html

# import packages
import torch
import torchvision

# Device configuration.
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

# Hyper-parameters
num_epochs = 10
num_classes = 10
batch_size = 100
learning_rate = 0.001
momentum = 0.9

# Load downloaded dataset.
import numpy as np
import gzip
import os
class MNISTDataset(torch.utils.data.Dataset):
    def __init__(self, root, train=True, transform=None):
        self.file_pre = 'train' if train == True else 't10k'
        self.transform = transform
        self.label_path = os.path.join(root, '%s-labels-idx1-ubyte.gz' % self.file_pre)
        self.image_path = os.path.join(root, '%s-images-idx3-ubyte.gz' % self.file_pre)
        self.images, self.labels = self.__read_data__(self.image_path, self.label_path)
    
    def __read_data__(self, image_path, label_path):
        # Read dataset.
        with gzip.open(label_path, 'rb') as lbpath:
            labels = np.frombuffer(lbpath.read(), np.uint8, offset=8)
        with gzip.open(image_path, 'rb') as imgpath:
            images = np.frombuffer(imgpath.read(), np.uint8, offset=16).reshape(len(labels), 28, 28)
        return images, labels
    
    def __getitem__(self, index):
        image, label = self.images[index], int