欢迎访问宙启技术站
智能推送

用Python编写的LeNet模型网络架构生成器

发布时间:2023-12-11 06:26:23

LeNet是一种经典的深度学习模型,由Yann LeCun等人于1998年提出,主要应用于手写数字识别。LeNet模型具有简单、有效的网络架构,在很多任务中都能够取得不错的性能。

在Python中,可以使用深度学习框架TensorFlow或者PyTorch来实现LeNet模型。下面我将介绍如何使用TensorFlow和PyTorch分别编写LeNet模型,并给出简单的使用例子。

使用TensorFlow实现LeNet模型:

import tensorflow as tf

def lenet(input_shape, num_classes):
    model = tf.keras.models.Sequential()
    model.add(tf.keras.layers.Conv2D(6, kernel_size=(5, 5), activation='relu', input_shape=input_shape))
    model.add(tf.keras.layers.MaxPooling2D(pool_size=(2, 2)))
    model.add(tf.keras.layers.Conv2D(16, kernel_size=(5, 5), activation='relu'))
    model.add(tf.keras.layers.MaxPooling2D(pool_size=(2, 2)))
    model.add(tf.keras.layers.Flatten())
    model.add(tf.keras.layers.Dense(120, activation='relu'))
    model.add(tf.keras.layers.Dense(84, activation='relu'))
    model.add(tf.keras.layers.Dense(num_classes, activation='softmax'))
    
    return model

# 使用例子:
input_shape = (32, 32, 1)  # 输入图像大小
num_classes = 10  # 类别数

model = lenet(input_shape, num_classes)
model.compile(optimizer='adam', loss='categorical_crossentropy', metrics=['accuracy'])

# 加载数据集
(x_train, y_train), (x_test, y_test) = tf.keras.datasets.mnist.load_data()
x_train = x_train.reshape(-1, 28, 28, 1)
x_test = x_test.reshape(-1, 28, 28, 1)
y_train = tf.keras.utils.to_categorical(y_train, num_classes)
y_test = tf.keras.utils.to_categorical(y_test, num_classes)

# 训练模型
model.fit(x_train, y_train, batch_size=128, epochs=10, validation_data=(x_test, y_test))

使用PyTorch实现LeNet模型:

import torch
import torch.nn as nn

class LeNet(nn.Module):
    def __init__(self, num_classes=10):
        super(LeNet, self).__init__()
        self.conv1 = nn.Conv2d(1, 6, kernel_size=5)
        self.pool = nn.MaxPool2d(kernel_size=2, stride=2)
        self.conv2 = nn.Conv2d(6, 16, kernel_size=5)
        self.fc1 = nn.Linear(16*4*4, 120)
        self.fc2 = nn.Linear(120, 84)
        self.fc3 = nn.Linear(84, num_classes)
        
    def forward(self, x):
        x = self.pool(torch.relu(self.conv1(x)))
        x = self.pool(torch.relu(self.conv2(x)))
        x = x.view(-1, 16*4*4)
        x = torch.relu(self.fc1(x))
        x = torch.relu(self.fc2(x))
        x = self.fc3(x)
        return x

# 使用例子:
input_shape = (1, 32, 32)  # 输入图像大小
num_classes = 10  # 类别数

model = LeNet(num_classes)
criterion = nn.CrossEntropyLoss()
optimizer = torch.optim.Adam(model.parameters())

# 加载数据集
transform = transforms.Compose([
    transforms.ToTensor(),
    transforms.Normalize((0.1307,), (0.3081,))
])
trainset = torchvision.datasets.MNIST(root='./data', train=True, download=True, transform=transform)
trainloader = torch.utils.data.DataLoader(trainset, batch_size=128, shuffle=True)

# 训练模型
for epoch in range(10):
    for inputs, labels in trainloader:
        optimizer.zero_grad()
        output = model(inputs)
        loss = criterion(output, labels)
        loss.backward()
        optimizer.step()

以上给出了在TensorFlow和PyTorch中使用LeNet模型的示例代码。你可以根据自己的需要进行相应的修改,例如更改输入图像的大小、类别数等。同时,你可以根据具体的任务来调整网络的超参数,例如卷积核大小、池化层大小等。希望以上代码能帮助你理解LeNet模型的网络架构和使用方法。