#AI挑战营第一站#PyTorch+MNIST数据集手写数字识别

天空中

#AI挑战营第一站#PyTorch+MNIST数据集手写数字识别 [复制链接]

# 模型训练环境
CPU：AMDR7
PyTorch版本：2.0
# 模型训练及其流程
## 载入头文件 
```
import torch
import numpy as np
from matplotlib import pyplot as plt
from torch.utils.data import DataLoader
from torchvision import transforms
from torchvision import datasets
import torch.nn.functional as F
```
## 定义超参数
```
Batch_Size = 64
Learning_Rate = 0.01
Momentum = 0.5
Period = 10
```
## 载入数据集
```
# 定义数据转换，将图像转换为张量
transform = transforms.Compose([transforms.ToTensor(), transforms.Normalize((0.1307,), (0.3081,))])
# 加载训练集和测试集，如果数据集不存在则自动下载
train_dataset = datasets.MNIST(root='./mnist', train=True, download=True, transform=transform)  
test_dataset = datasets.MNIST(root='./mnist', train=False, download=True, transform=transform)
```
## 加载数据集
```
train_loader = DataLoader(train_dataset, batch_size=Batch_Size, shuffle=True)
test_loader = DataLoader(test_dataset, batch_size=Batch_Size, shuffle=False)
```
## 输出数据集图像
```
fig = plt.figure()
for i in range(12):
    plt.subplot(3, 4, i+1)
    plt.tight_layout()
    plt.imshow(train_dataset.train_data<i>, cmap='gray', interpolation='none')
    plt.title("Labels: {}".format(train_dataset.train_labels<i>))
    plt.xticks([])
    plt.yticks([])
plt.show()
```

<ignore_js_op>

<a href="javascript:;" onclick="showWindow(this.id, this.getAttribute('url'), 'get', 0);" id="savephoto_802891" url="home.php?mod=spacecp&ac=album&op=saveforumphoto&aid=802891&handlekey=savephoto_802891">保存到相册</a>

</p>

</div>
<div class="tip_horn"></div>
</div>

</ignore_js_op>

## 定义模型
```
import torch

class MODEL_NUM(torch.nn.Module):
    def __init__(self):
        super(Net, self).__init__()
        # 第一个卷积层
        self.conv1 = torch.nn.Sequential(
            torch.nn.Conv2d(1, 10, kernel_size=5),  # 输入通道为1，输出通道为10，卷积核大小为5x5
            torch.nn.ReLU(),  # ReLU 激活函数
            torch.nn.MaxPool2d(kernel_size=2),  # 最大池化层，池化窗口大小为2x2
        )
        # 第二个卷积层
        self.conv2 = torch.nn.Sequential(
            torch.nn.Conv2d(10, 20, kernel_size=5),  # 输入通道为10，输出通道为20，卷积核大小为5x5
            torch.nn.ReLU(),  # ReLU 激活函数
            torch.nn.MaxPool2d(kernel_size=2),  # 最大池化层，池化窗口大小为2x2
        )
        # 全连接层
        self.fc = torch.nn.Sequential(
            torch.nn.Linear(320, 50),  # 输入大小为320，输出大小为50
            torch.nn.Linear(50, 10),    # 输入大小为50，输出大小为10（对应0~9十个类别）
        )

def forward(self, x):
        Batch_Size = x.size(0)  # 获取输入张量的批量大小
        x = self.conv1(x)  # 进行第一个卷积层的前向传播
        x = self.conv2(x)  # 进行第二个卷积层的前向传播
        x = x.view(Batch_Size, -1)  # 将张量展平为一维，以便进入全连接层
        x = self.fc(x)  # 进行全连接层的前向传播
        return x  # 返回模型的输出，维度为10，对应0~9十个类别

```
## 定义损失函数和优化器
```
criterion = torch.nn.CrossEntropyLoss()		#损失函数
# 随机梯度下降 (SGD) 优化器，lr 参数指定学习率，momentum 参数指定动量
optimizer = torch.optim.SGD(model.parameters(), lr=Learning_Rate, momentum=Momentum)

```
## 定义训练函数和测试函数
```
def train(epoch):
    running_loss = 0.0  # 用于累计每个 epoch 的总损失
    running_total = 0  # 用于累计每个 epoch 的总样本数
    running_correct = 0  # 用于累计每个 epoch 的总正确预测数
    for batch_idx, data in enumerate(train_loader, 0):
        inputs, target = data
        optimizer.zero_grad()  # 梯度清零
        outputs = model(inputs)  # 前向传播
        loss = criterion(outputs, target)  # 计算损失
        loss.backward()  # 反向传播
        optimizer.step()  # 更新参数
        running_loss += loss.item()  # 累加损失

predicted = torch.max(outputs.data, dim=1)[1]  # 获取模型预测结果中的类别索引
        running_total += inputs.shape[0]  # 累加样本数
        running_correct += (predicted == target).sum().item()  # 累加正确预测数

if batch_idx % 300 == 299:  # 每300个 mini-batch 输出一次损失和准确率
            print('[%d, %5d]: loss: %.3f , acc: %.2f %%' % (epoch + 1, batch_idx + 1, running_loss / 300, 100 * running_correct / running_total))
            running_loss = 0.0  # 清零损失
            running_total = 0  # 清零样本数
            running_correct = 0  # 清零正确预测数
```
```
def test():
    correct = 0  # 用于累计测试集中的正确预测数
    total = 0  # 用于累计测试集中的总样本数
    with torch.no_grad():  # 关闭梯度计算，因为在测试阶段不需要计算梯度
        for data in test_loader:  # 遍历测试集的每个 mini-batch
            images, labels = data  # 获取输入图像和标签
            outputs = model(images)  # 模型预测
            _, predicted = torch.max(outputs.data, dim=1)  # 获取预测结果中的类别索引
            total += labels.size(0)  # 累加样本数
            correct += (predicted == labels).sum().item()  # 累加正确预测数

acc = correct / total  # 计算测试集准确率
    print('[%d / %d]: Accuracy on test set: %.1f %% ' % (epoch+1, Period, 100 * acc))  # 打印测试集准确率
    return acc  # 返回测试集准确率

```
## 主函数
```
if __name__ == '__main__':
    acc_list_test = []
    for epoch in range(Period):
        train(epoch)
        if epoch % 10 == 9:  #每训练10轮 测试1次,训练次数少于10轮是将其注释
        	acc_test = test()
        	acc_list_test.append(acc_test)

torch.save(model.state_dict(), './MODE_NUM.pth')

plt.plot(acc_list_test)
    plt.xlabel('Period')
    plt.ylabel('Accuracy On TestSet')
    plt.show()
```
## 训练结果展示

<ignore_js_op>

<a href="javascript:;" onclick="showWindow(this.id, this.getAttribute('url'), 'get', 0);" id="savephoto_802892" url="home.php?mod=spacecp&ac=album&op=saveforumphoto&aid=802892&handlekey=savephoto_802892">保存到相册</a>

</p>

</div>
<div class="tip_horn"></div>
</div>

</ignore_js_op>

<ignore_js_op>

<a href="javascript:;" onclick="showWindow(this.id, this.getAttribute('url'), 'get', 0);" id="savephoto_802893" url="home.php?mod=spacecp&ac=album&op=saveforumphoto&aid=802893&handlekey=savephoto_802893">保存到相册</a>

</p>

</div>
<div class="tip_horn"></div>
</div>

</ignore_js_op>

## 模型转换
```
# 创建模型实例并加载预训练的模型参数
model = MODEL_NUM()
model.load_state_dict(torch.load('./MODE_NUM.pth'))  # 加载预训练的模型参数
model.eval()  # 将模型设置为评估模式
# 设置输入张量的形状（batch_size, channels, height, width）
dummy_input = torch.randn(1, 1, 28, 28)  # 1个样本，1个通道，28x28图像
# 导出模型
torch.onnx.export(model, dummy_input, "./NUM_ONNX.onnx", verbose=True)
```
## 训练模型分享

<ignore_js_op>

<em class="xg1">(88.39 KB, 下载次数: 1)</em>
</span>
<div class="tip tip_4" id="attach_802890_menu" style="position: absolute; display: none" disautofocus="true">
<div class="tip_c xs0">
<div class="y">2024-4-17 18:33 上传</div>
点击文件名下载附件

</div>
<div class="tip_horn"></div>
</div>

</ignore_js_op>

<ignore_js_op>

<em class="xg1">(124.67 KB, 下载次数: 1)</em>
</span>
<div class="tip tip_4" id="attach_802889_menu" style="position: absolute; display: none" disautofocus="true">
<div class="tip_c xs0">
<div class="y">2024-4-17 18:33 上传</div>
点击文件名下载附件

</div>
<div class="tip_horn"></div>
</div>

</ignore_js_op>

Jacktang

挑战营第一站#PyTorch+MNIST数据集手写数字识别，整的漂亮

wangerxian

不错不错，流程很清晰，排版也很好看~

#AI挑战营第一站#PyTorch+MNIST数据集手写数字识别 [复制链接]

最新回复