机器学习之卷积神经网络Lenet5训练模型

最新推荐文章于 2025-05-13 23:38:20 发布

原创

最新推荐文章于 2025-05-13 23:38:20 发布 · 466 阅读

2 ·

CC 4.0 BY-SA版权

文章标签：

#cnn #机器学习 #pytorch

Lenet5训练模型

下载数据集

在这里插入图片描述
可以提前下载也可以在线下载

train_data = torchvision.datasets.MNIST(root='./',download=True,train=True,transform=transform)
test_data = torchvision.datasets.MNIST(root='./',download=True,train=False,transform=transform)

训练模型

import torch
import torchvision
class Lenet5(torch.nn.Module):
    def __init__(self):
        super(Lenet5, self).__init__()
        self.model = torch.nn.Sequential(
            torch.nn.Conv2d(in_channels=1,out_channels=6,kernel_size=5), # 1*32*32 # 6*28*28
            torch.nn.ReLU(),
            torch.nn.MaxPool2d(kernel_size=2,  # 6*14*14
                                stride=2),
            torch.nn.Conv2d(in_channels=6,
                                    out_channels=16,
                                    kernel_size=5),  # 16 *10*10
            torch.nn.ReLU(),
            torch.nn.MaxPool2d(kernel_size=2,
                                stride=2),  # 16*5*5
            to