浏览器中加载CNN进行手写数字识别，并部署到Gitee Page

2023-09-10 03:29:13

准备工作

1.1 安装依赖

pip install torch torchvision numpy matplotlib

1.2 下载并处理数据

from torchvision import datasets, transforms

# 下载Mnist数据集
train_data = datasets.MNIST(
    root='./data',
    train=True,
    download=True,
    transform=transforms.ToTensor()
)

# 下载测试集
test_data = datasets.MNIST(
    root='./data',
    train=False,
    download=True,
    transform=transforms.ToTensor()
)

2. 训练CNN网络

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

# 定义CNN网络结构
class CNN(nn.Module):
    def __init__(self):
        super(CNN, self).__init__()
        self.conv1 = nn.Conv2d(1, 32, 3, 1)
        self.pool1 = nn.MaxPool2d(2, 2)
        self.conv2 = nn.Conv2d(32, 64, 3, 1)
        self.pool2 = nn.MaxPool2d(2, 2)
        self.fc1 = nn.Linear(64 * 4 * 4, 128)
        self.fc2 = nn.Linear(128, 10)

    def forward(self, x):
        x = self.conv1(x)
        x = self.pool1(x)
        x = self.conv2(x)
        x = self.pool2(x)
        x = x.view(x.size(0), -1)
        x = self.fc1(x)
        x = self.fc2(x)
        return x

# 实例化CNN网络
model = CNN()

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

# 训练网络
for epoch in range(10):
    for i, (inputs, labels) in enumerate(train_data):
        # 正向传播
        outputs = model(inputs)

        # 计算损失
        loss = criterion(outputs, labels)

        # 反向传播
        optimizer.zero_grad()
        loss.backward()

        # 更新参数
        optimizer.step()

# 保存训练好的模型
torch.save(model.state_dict(), 'mnist_cnn.pt')

3. 搭建静态页面

<!DOCTYPE html>
<html>
<head>
  
  <script src="mnist_cnn.pt"></script>
  <script>
    // 加载模型
    const model = torch.load('mnist_cnn.pt');

    // 获取canvas元素
    const canvas = document.getElementById('canvas');
    const ctx = canvas.getContext('2d');

    // 添加事件监听器，在canvas上绘制数字时触发
    canvas.addEventListener('mousemove', function(e) {
      if (e.buttons !== 1) {
        return;
      }

      // 获取鼠标坐标
      const x = e.clientX - canvas.offsetLeft;
      const y = e.clientY - canvas.offsetTop;

      // 在canvas上绘制线条
      ctx.lineTo(x, y);
      ctx.stroke();
    });

    // 添加按钮，点击后识别canvas上的数字
    document.getElementById('predict-button').addEventListener('click', function() {
      // 获取canvas上的图像数据
      const imageData = ctx.getImageData(0, 0, 28, 28);

      // 将图像数据转换为张量
      const tensor = torch.from_numpy(imageData.data).view(1, 1, 28, 28);

      // 预测数字
      const output = model(tensor);

      // 获取预测结果
      const predicted = output.argmax(dim=1).item();

      // 显示预测结果
      document.getElementById('result').innerHTML = predicted;
    });
  </script>
</head>
<body>
  <h1>手写数字识别</h1>
  <canvas id="canvas" width="280" height="280"></canvas>
  <button id="predict-button">识别</button>
  <div id="result"></div>
</body>
</html>