#AI挑战营第一站# PC基于PyTorch的MNIST模型训练过程与模型转换

knv

#AI挑战营第一站# PC基于PyTorch的MNIST模型训练过程与模型转换 [复制链接]

很高兴参加这个活动，我使用miniconda进行环境配置，并使用pycharm进行开发，电脑配置是 i5 8300h+ 1050ti.本次训练使用PyTorch，与MNIST数据集。

安装完成后，使用conda命令激活虚拟环境。并使用pycharm进行开发

Conda安装pytorch环境：

conda install pytorch torchvision torchaudio pytorch-cuda=11.8 -c pytorch -c nvidia

定义了一个模型，包含三个卷积层和两个全连接层。每个卷积层后面都跟着一个批量归一化(Batch Normalization)层和ReLU激活函数，然后是一个最大池化层。在全连接层之间，还有一个Dropout层用于防止过拟合。

class OptimizedConvNet(nn.Module):
    def __init__(self, num_classes=10):
        super(OptimizedConvNet, self).__init__()
        self.layer1 = nn.Sequential(
            nn.Conv2d(1, 16, kernel_size=5, stride=1, padding=2),
            nn.BatchNorm2d(16),
            nn.ReLU(),
            nn.MaxPool2d(kernel_size=2, stride=2))
        self.layer2 = nn.Sequential(
            nn.Conv2d(16, 32, kernel_size=5, stride=1, padding=2),
            nn.BatchNorm2d(32),
            nn.ReLU(),
            nn.MaxPool2d(kernel_size=2, stride=2))
        self.layer3 = nn.Sequential(
            nn.Conv2d(32, 64, kernel_size=5, stride=1, padding=2),
            nn.BatchNorm2d(64),
            nn.ReLU(),
            nn.MaxPool2d(kernel_size=2, stride=2))
        self.fc1 = nn.Linear(3 * 3 * 64, 500)
        self.dropout = nn.Dropout(0.5)
        self.fc2 = nn.Linear(500, num_classes)

    def forward(self, x):
        out = self.layer1(x)
        out = self.layer2(out)
        out = self.layer3(out)
        out = out.reshape(out.size(0), -1)
        out = self.fc1(out)
        out = self.dropout(out)
        out = self.fc2(out)
        return out

完整代码：（使用GPU进行训练）

import torch
import torch.nn as nn
import torch.optim as optim
from torchvision import datasets, transforms
from torch.utils.data import DataLoader


class OptimizedConvNet(nn.Module):
    def __init__(self, num_classes=10):
        super(OptimizedConvNet, self).__init__()
        self.layer1 = nn.Sequential(
            nn.Conv2d(1, 16, kernel_size=5, stride=1, padding=2),
            nn.BatchNorm2d(16),
            nn.ReLU(),
            nn.MaxPool2d(kernel_size=2, stride=2))
        self.layer2 = nn.Sequential(
            nn.Conv2d(16, 32, kernel_size=5, stride=1, padding=2),
            nn.BatchNorm2d(32),
            nn.ReLU(),
            nn.MaxPool2d(kernel_size=2, stride=2))
        self.layer3 = nn.Sequential(
            nn.Conv2d(32, 64, kernel_size=5, stride=1, padding=2),
            nn.BatchNorm2d(64),
            nn.ReLU(),
            nn.MaxPool2d(kernel_size=2, stride=2))
        self.fc1 = nn.Linear(3 * 3 * 64, 500)
        self.dropout = nn.Dropout(0.5)
        self.fc2 = nn.Linear(500, num_classes)

    def forward(self, x):
        out = self.layer1(x)
        out = self.layer2(out)
        out = self.layer3(out)
        out = out.reshape(out.size(0), -1)
        out = self.fc1(out)
        out = self.dropout(out)
        out = self.fc2(out)
        return out

# 定义超参数
learning_rate = 0.001
batch_size = 64
num_epochs = 10


# 检查是否有可用的GPU
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

# 创建模型，并将模型移动到GPU上
model = OptimizedConvNet(num_classes=10).to(device)

# 定义损失函数和优化器
criterion = nn.CrossEntropyLoss()
optimizer = torch.optim.Adam(model.parameters(), lr=learning_rate)

# 加载MNIST数据集
transform = transforms.Compose([transforms.ToTensor(), transforms.Normalize((0.1307,), (0.3081,))])
train_dataset = datasets.MNIST(root='./data', train=True, download=True, transform=transform)
train_loader = DataLoader(train_dataset, batch_size=64, shuffle=True)
# 加载验证集
valid_dataset = datasets.MNIST(root='./data', train=False, download=True, transform=transform)
valid_loader = DataLoader(valid_dataset, batch_size=64, shuffle=False)

# 初始化最优验证损失为无穷大
best_valid_loss = float('inf')

# 训练模型
for epoch in range(num_epochs):
    # 在训练集上训练
    model.train()
    total_train_loss = 0
    for i, (images, labels) in enumerate(train_loader):
        images = images.to(device)
        labels = labels.to(device)

        outputs = model(images)
        loss = criterion(outputs, labels)

        optimizer.zero_grad()
        loss.backward()
        optimizer.step()
        if (i + 1) % 100 == 0:
            print(f'Epoch [{epoch + 1}/10], Step [{i + 1}/{len(train_loader)}], Loss: {loss.item()}')

        total_train_loss += loss.item()

    average_train_loss = total_train_loss / len(train_loader)
    print(f'Epoch: {epoch + 1}, Training Loss: {average_train_loss:.4f}')

    # 在验证集上验证
    model.eval()
    valid_loss = 0.0
    with torch.no_grad():
        for images, labels in valid_loader:
            images = images.to(device)
            labels = labels.to(device)

            outputs = model(images)
            loss = criterion(outputs, labels)
            valid_loss += loss.item()

    valid_loss /= len(valid_loader)
    print(f'Epoch: {epoch+1}, Validation Loss: {valid_loss:.4f}')

    # 如果验证损失有所下降，则保存模型
    if valid_loss < best_valid_loss:
        best_valid_loss = valid_loss
        torch.save(model.state_dict(), 'best_model.pth')
        print('Model saved.')

print("Training completed.")

训练过程：

在代码中，每次训练完毕后，保存最佳模型。

模型转ONNX

安装ONNX相关依赖库，运行代码即可转换完毕。

conda install -c conda-forge onnx

import torch
import torch.nn as nn
import torchvision

class OptimizedConvNet(nn.Module):
    def __init__(self, num_classes=10):
        super(OptimizedConvNet, self).__init__()
        self.layer1 = nn.Sequential(
            nn.Conv2d(1, 16, kernel_size=5, stride=1, padding=2),
            nn.BatchNorm2d(16),
            nn.ReLU(),
            nn.MaxPool2d(kernel_size=2, stride=2))
        self.layer2 = nn.Sequential(
            nn.Conv2d(16, 32, kernel_size=5, stride=1, padding=2),
            nn.BatchNorm2d(32),
            nn.ReLU(),
            nn.MaxPool2d(kernel_size=2, stride=2))
        self.layer3 = nn.Sequential(
            nn.Conv2d(32, 64, kernel_size=5, stride=1, padding=2),
            nn.BatchNorm2d(64),
            nn.ReLU(),
            nn.MaxPool2d(kernel_size=2, stride=2))
        self.fc1 = nn.Linear(3 * 3 * 64, 500)
        self.dropout = nn.Dropout(0.5)
        self.fc2 = nn.Linear(500, num_classes)

    def forward(self, x):
        out = self.layer1(x)
        out = self.layer2(out)
        out = self.layer3(out)
        out = out.reshape(out.size(0), -1)
        out = self.fc1(out)
        out = self.dropout(out)
        out = self.fc2(out)
        return out

# 加载模型
model = OptimizedConvNet()
model.load_state_dict(torch.load('best_model.pth'))
model.eval()

# 创建一个模拟输入，这需要和您的模型输入维度相同
# 例如，如果您的模型接收1x28x28的图像作为输入，您可以创建一个相应大小的随机张量
dummy_input = torch.randn(1, 1, 28, 28)

# 导出模型
torch.onnx.export(model, dummy_input, "model.onnx")

train.py (3.67 KB, 下载次数: 1)

toOnnx.py (1.62 KB, 下载次数: 1)

model.onnx (1.37 MB, 下载次数: 1)

best_model.pth (1.38 MB, 下载次数: 1)

补充内容 (2024-5-9 08:56): 准确率补充到楼下了

Jacktang

电脑配置是 i5 8300h+ 1050ti.玩这个训练，用PyTorch，与MNIST数据集应该没有问题

knv

本帖最后由 knv 于 2024-5-8 18:10 编辑

补充准确率：验证准确率 99.16%

新版本代码：

import torch
import torch.nn as nn
import torch.optim as optim
from torchvision import datasets, transforms
from torch.utils.data import DataLoader


class OptimizedConvNet(nn.Module):
    def __init__(self,num_classes=10):
        super(OptimizedConvNet, self).__init__()
        self.layer1 = nn.Sequential(
            nn.Conv2d(1, 32, kernel_size=5, stride=1, padding=2),
            nn.ReLU(),
            nn.MaxPool2d(kernel_size=2, stride=2))
        self.layer2 = nn.Sequential(
            nn.Conv2d(32, 64, kernel_size=5, stride=1, padding=2),
            nn.ReLU(),
            nn.MaxPool2d(kernel_size=2, stride=2))
        self.drop_out = nn.Dropout()
        self.fc1 = nn.Linear(7 * 7 * 64, 1000)
        self.fc2 = nn.Linear(1000, num_classes)

    def forward(self, x):
        out = self.layer1(x)
        out = self.layer2(out)
        out = out.reshape(out.size(0), -1)
        out = self.drop_out(out)
        out = self.fc1(out)
        out = self.fc2(out)
        return out

# 定义超参数
learning_rate = 0.001
batch_size = 64
num_epochs = 10


# 检查是否有可用的GPU
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

# 创建模型，并将模型移动到GPU上
model = OptimizedConvNet(num_classes=10).to(device)

# 定义损失函数和优化器
criterion = nn.CrossEntropyLoss()
optimizer = torch.optim.Adam(model.parameters(), lr=learning_rate)

# 加载MNIST数据集
transform = transforms.Compose([transforms.ToTensor(), transforms.Normalize((0.1307,), (0.3081,))])
train_dataset = datasets.MNIST(root='./data', train=True, download=True, transform=transform)
train_loader = DataLoader(train_dataset, batch_size=64, shuffle=True)
# 加载验证集
valid_dataset = datasets.MNIST(root='./data', train=False, download=True, transform=transform)
valid_loader = DataLoader(valid_dataset, batch_size=64, shuffle=False)

# 初始化最优验证损失为无穷大
best_valid_loss = float('inf')
import matplotlib.pyplot as plt

# 初始化准确率列表
train_accuracies = []
valid_accuracies = []

for epoch in range(num_epochs):
    # 在训练集上训练
    model.train()
    total_train_loss = 0
    correct_train_preds = 0
    total_train_preds = 0

    for i, (images, labels) in enumerate(train_loader):
        images = images.to(device)
        labels = labels.to(device)

        outputs = model(images)
        loss = criterion(outputs, labels)

        _, predicted = torch.max(outputs.data, 1)
        total_train_preds += labels.size(0)
        correct_train_preds += (predicted == labels).sum().item()

        optimizer.zero_grad()
        loss.backward()
        optimizer.step()

        if (i + 1) % 100 == 0:
            print(f'Epoch [{epoch + 1}/{num_epochs}], Step [{i + 1}/{len(train_loader)}], Loss: {loss.item()}')

        total_train_loss += loss.item()

    average_train_loss = total_train_loss / len(train_loader)
    train_accuracy = correct_train_preds / total_train_preds
    train_accuracies.append(train_accuracy)
    print(f'Epoch: {epoch + 1}, Training Loss: {average_train_loss:.4f}, Training Accuracy: {train_accuracy:.4f}')

    # 在验证集上验证
    model.eval()
    valid_loss = 0.0
    correct_valid_preds = 0
    total_valid_preds = 0
    with torch.no_grad():
        for images, labels in valid_loader:
            images = images.to(device)
            labels = labels.to(device)

            outputs = model(images)
            loss = criterion(outputs, labels)

            _, predicted = torch.max(outputs.data, 1)
            total_valid_preds += labels.size(0)
            correct_valid_preds += (predicted == labels).sum().item()

            valid_loss += loss.item()

    valid_loss /= len(valid_loader)
    valid_accuracy = correct_valid_preds / total_valid_preds
    valid_accuracies.append(valid_accuracy)
    print(f'Epoch: {epoch+1}, Validation Loss: {valid_loss:.4f}, Validation Accuracy: {valid_accuracy:.4f}')

    # 如果验证损失有所下降，则保存模型
    if valid_loss < best_valid_loss:
        best_valid_loss = valid_loss
        torch.save(model.state_dict(), 'best_model.pth')
        print('Model saved.')

# 绘制训练和验证准确率
plt.plot(range(1, num_epochs + 1), train_accuracies, label='Train')
plt.plot(range(1, num_epochs + 1), valid_accuracies, label='Valid')
plt.title('Model Accuracy')
plt.xlabel('Epochs')
plt.ylabel('Accuracy')
plt.legend()
plt.show()

#AI挑战营第一站# PC基于PyTorch的MNIST模型训练过程与模型转换 [复制链接]

最新回复