#AI Challenge Camp First Stop#Complete handwritten digital model training on PC

Fardreamer

#AI Challenge Camp First Stop#Complete handwritten digital model training on PC [Copy link]

model.py

import torch.nn as nn
import torch.nn.functional as F


class Net(nn.Module):
    class Net(nn.Module):
        """
        构造一个具有三个全连接层的神经网络模型。
        
        参数:
        无
        
        返回值:
        无
        """
        def __init__(self):
            super(Net, self).__init__()
            self.fc1 = nn.Linear(28 * 28, 512)  # 第一层全连接层，输入维度为28*28，输出维度为512
            self.fc2 = nn.Linear(512, 256)  # 第二层全连接层，输入维度为512，输出维度为256
            self.fc3 = nn.Linear(256, 10)  # 第三层全连接层，输入维度为256，输出维度为10，用于输出类别概率

    def forward(self, x):
        """
        前向传播函数：将输入的图像数据通过一个标准的卷积神经网络结构进行处理，得到分类的log softmax输出。
    
        参数:
        - x : 输入的图像张量，形状为(batch_size, 1, 28, 28)，其中batch_size是批次大小。
    
        返回值:
        - 经过卷积神经网络处理后，得到的log softmax输出张量，形状为(batch_size, 10)，其中10是类别数。
        """
        # 将输入图像张量展平为一维
        x = x.view(-1, 28 * 28)
        # 通过第一个全连接层并应用ReLU激活函数
        x = F.relu(self.fc1(x))
        # 通过第二个全连接层并应用ReLU激活函数
        x = F.relu(self.fc2(x))
        # 通过第三个全连接层，此层通常作为输出层，不应用激活函数
        x = self.fc3(x)
        # 应用log softmax激活函数，以得到类别概率的对数
        return F.log_softmax(x, dim=1)

train.py

import torch
import torch.nn as nn
import torch.optim as optim


def train_model(model, device, train_loader, epochs=2, learning_rate=0.001):
    """
    训练给定的模型。
    
    参数:
    - model: 要训练的模型。
    - device: 指定模型运行的设备（如"cpu"或"cuda"）。
    - train_loader: 训练数据的加载器。
    - epochs: 训练的轮数，默认为2。
    - learning_rate: 学习率，默认为0.001。
    
    返回值:
    - 无
    """
    criterion = nn.NLLLoss()  # 定义损失函数
    optimizer = optim.Adam(model.parameters(), lr=learning_rate)  # 使用Adam优化器

    for epoch in range(epochs):  # 遍历训练轮数
        running_loss = 0.0  # 初始化运行时损失
        for i, data in enumerate(train_loader, 0):  # 遍历训练数据集
            inputs, labels = data[0].to(device), data[1].to(device)  # 将数据移至指定设备
            optimizer.zero_grad()  # 清除之前的梯度
            outputs = model(inputs)  # 通过模型获取输出
            loss = criterion(outputs, labels)  # 计算损失
            loss.backward()  # 反向传播
            optimizer.step()  # 更新模型参数
            running_loss += loss.item()  # 累加损失
            if i % 2000 == 1999:  # 每2000个样本打印一次损失
                print('[%d, %5d] loss: %.3f' % (epoch + 1, i + 1, running_loss / 2000))
                running_loss = 0.0
    print('Finished Training')  # 训练完成

test.py

import torch


def test_model(model, device, test_loader):
    """
    测试模型的准确率。

    参数:
        model (torch.nn.Module): 要测试的模型
        device (str or torch.device): 设备，如'cuda'或'cpu'
        test_loader (torch.utils.data.DataLoader): 测试数据集加载器

    返回:
        float: 模型的准确率
    """
    model.eval()  # 将模型设置为评估模式
    correct = 0
    total = 0

    with torch.no_grad():  # 不计算梯度，节省内存和计算资源
        for images, labels in test_loader:
            images, labels = images.to(device), labels.to(device)  # 将数据和标签移动到指定设备
            outputs = model(images)  # 获取模型输出
            _, predicted = torch.max(outputs, 1)  # 获取预测类别
            total += labels.size(0)  # 更新总数
            correct += (predicted == labels).sum().item()  # 更新正确预测数

    # 计算准确率
    accuracy = 100 * correct / total
    print(f'Accuracy of the network on the test images: {accuracy:.2f} %')
    return accuracy  # 返回准确率，以便在其他地方使用

main.py

import torch
import torchvision
import torchvision.transforms as transforms
from model import Net
from train import train_model
from test import test_model
# import onnx

# 数据预处理
transform = transforms.Compose(
    [transforms.ToTensor(),
     transforms.Normalize((0.5,), (0.5,))])

# 下载并加载数据集
trainset = torchvision.datasets.MNIST('~/.pytorch/MNIST_data/', download=True, train=True, transform=transform)
trainloader = torch.utils.data.DataLoader(trainset, batch_size=64, shuffle=True)

testset = torchvision.datasets.MNIST('~/.pytorch/MNIST_data/', download=True, train=False, transform=transform)
testloader = torch.utils.data.DataLoader(testset, batch_size=64, shuffle=False)

# 实例化模型
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
model = Net().to(device)

# 训练模型
train_model(model, device, trainloader)

# 测试模型
# test_model(model, device, testloader)

torch.save(model.state_dict(), 'model_weights.pth')

# 准备虚拟输入用于 ONNX 转换
# 假设输入是一个 1x1x28x28 的张量，因为 MNIST 数据集的图片大小是 28x28
dummy_input = torch.randn(1, 1, 28, 28, device=device)

# 确保模型处于评估模式
model.eval()

# 将模型转换为 ONNX 格式
onnx_file_path = 'model.onnx'
torch.onnx.export(model, dummy_input, onnx_file_path, opset_version=12)

print(f"Model exported to {onnx_file_path}")

model_weights.pth (2.05 MB, downloads: 0)

2024-4-20 23:38 上传

点击文件名下载附件

model.onnx (2.04 MB, downloads: 0)

2024-4-20 23:38 上传

点击文件名下载附件

Jacktang

Complete handwritten digital model training on PC, it sounds cool, model training is quite popular now

nian1206

Mark it and I will use it later. Thank you very much for sharing. May good people have a peaceful life!