{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 仅全链接层实现手写数字识别,tensorflow版本\n",
    "\n",
    "## 模型定义\n",
    "\n",
    "- 输入层：28*28=784\n",
    "- 输出层：10\n",
    "- 隐藏层：784(input)->256(relu)(dropout-20%)->128(tanh)(dropout-20%)->10(softmax)\n",
    "\n",
    "## 数据处理\n",
    "\n",
    "- tf集成的keras中提供了mnist数据集，其数据格式为`numpy.ndarray`\n",
    "- 将数据集展平为一维向量并进行归一化\n",
    "- 将标签进行one-hot编码\n",
    "\n",
    "## 训练\n",
    "\n",
    "- 优化器为`Adam`\n",
    "- 损失函数为`categorical_crossentropy` - 分类交叉熵\n",
    "- 评估指标为`accuracy`\n",
    "- 模型训练时划分训练集和验证集，由验证集对模型进行评估以调整学习率\n",
    "- 保存模型到本地 - tf默认保存的内容包括模型的结构、权重、优化器的状态等\n",
    "\n",
    "## 测试\n",
    "\n",
    "- 使用keras的接口对模型进行加载\n",
    "- 对传入的图片进行处理，先要变28*28，再变一维向量，最后使用同样的归一化\n",
    "- 输出的结果是10个数字的概率，使用tf的argmax函数获取概率最大的数字即为预测结果\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-03-13T12:25:56.129548Z",
     "start_time": "2025-03-13T12:25:13.734410Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "可用设备： [PhysicalDevice(name='/physical_device:CPU:0', device_type='CPU'), PhysicalDevice(name='/physical_device:GPU:0', device_type='GPU')]\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "2025-03-14 18:29:18.749063: I metal_plugin/src/device/metal_device.cc:1154] Metal device set to: Apple M1\n",
      "2025-03-14 18:29:18.749282: I metal_plugin/src/device/metal_device.cc:296] systemMemory: 16.00 GB\n",
      "2025-03-14 18:29:18.749287: I metal_plugin/src/device/metal_device.cc:313] maxCacheSize: 5.33 GB\n",
      "2025-03-14 18:29:18.749596: I tensorflow/core/common_runtime/pluggable_device/pluggable_device_factory.cc:306] Could not identify NUMA node of platform GPU ID 0, defaulting to 0. Your kernel may not have been built with NUMA support.\n",
      "2025-03-14 18:29:18.750001: I tensorflow/core/common_runtime/pluggable_device/pluggable_device_factory.cc:272] Created TensorFlow device (/job:localhost/replica:0/task:0/device:GPU:0 with 0 MB memory) -> physical PluggableDevice (device: 0, name: METAL, pci bus id: <undefined>)\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 1/10\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "2025-03-14 18:29:19.283391: I tensorflow/core/grappler/optimizers/custom_graph_optimizer_registry.cc:117] Plugin optimizer for device_type GPU is enabled.\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "375/375 [==============================] - 6s 12ms/step - loss: 0.4132 - accuracy: 0.8765 - val_loss: 0.2169 - val_accuracy: 0.9364\n",
      "Epoch 2/10\n",
      "375/375 [==============================] - 4s 12ms/step - loss: 0.2421 - accuracy: 0.9277 - val_loss: 0.1686 - val_accuracy: 0.9514\n",
      "Epoch 3/10\n",
      "375/375 [==============================] - 5s 12ms/step - loss: 0.2034 - accuracy: 0.9388 - val_loss: 0.1450 - val_accuracy: 0.9585\n",
      "Epoch 4/10\n",
      "375/375 [==============================] - 5s 12ms/step - loss: 0.1796 - accuracy: 0.9459 - val_loss: 0.1336 - val_accuracy: 0.9612\n",
      "Epoch 5/10\n",
      "375/375 [==============================] - 5s 13ms/step - loss: 0.1649 - accuracy: 0.9498 - val_loss: 0.1226 - val_accuracy: 0.9633\n",
      "Epoch 6/10\n",
      "375/375 [==============================] - 4s 12ms/step - loss: 0.1542 - accuracy: 0.9530 - val_loss: 0.1183 - val_accuracy: 0.9644\n",
      "Epoch 7/10\n",
      "375/375 [==============================] - 4s 11ms/step - loss: 0.1479 - accuracy: 0.9549 - val_loss: 0.1147 - val_accuracy: 0.9664\n",
      "Epoch 8/10\n",
      "375/375 [==============================] - 4s 11ms/step - loss: 0.1399 - accuracy: 0.9571 - val_loss: 0.1139 - val_accuracy: 0.9665\n",
      "Epoch 9/10\n",
      "375/375 [==============================] - 4s 11ms/step - loss: 0.1357 - accuracy: 0.9584 - val_loss: 0.1085 - val_accuracy: 0.9672\n",
      "Epoch 10/10\n",
      "375/375 [==============================] - 4s 11ms/step - loss: 0.1340 - accuracy: 0.9580 - val_loss: 0.1044 - val_accuracy: 0.9693\n",
      "313/313 [==============================] - 3s 9ms/step - loss: 0.1019 - accuracy: 0.9700\n",
      "accuracy: 0.9700\n",
      "模型已保存到 ./model/mnist_model-tf.h5\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/Users/wolves/mambaforge/envs/ail/lib/python3.10/site-packages/keras/src/engine/training.py:3103: UserWarning: You are saving your model as an HDF5 file via `model.save()`. This file format is considered legacy. We recommend using instead the native Keras format, e.g. `model.save('my_model.keras')`.\n",
      "  saving_api.save_model(\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 57636 characters]",
      "text/plain": [
       "<Figure size 1200x400 with 2 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "import tensorflow as tf\n",
    "import tensorflow.keras as keras\n",
    "from keras import layers\n",
    "import matplotlib.pyplot as plt\n",
    "print(\"可用设备：\", tf.config.list_physical_devices())\n",
    "\n",
    "# 加载mnist数据集\n",
    "(x_train, y_train), (x_test, y_test) = keras.datasets.mnist.load_data()\n",
    "\n",
    "\n",
    "# 处理x数据，mnist数据集为灰度图片，范围为0-255，直接除以255，等同归一化\n",
    "x_train = x_train.reshape(-1, 28 * 28).astype('float32') / 255.0\n",
    "x_test = x_test.reshape(-1, 28 * 28).astype('float32') / 255.0\n",
    "\n",
    "# 处理y数据，mnist数据集为0-9的数字，需要将其转换为one-hot编码\n",
    "y_train = keras.utils.to_categorical(y_train,10)\n",
    "y_test = keras.utils.to_categorical(y_test,10)\n",
    "\n",
    "# 构建神经网络\n",
    "model = keras.Sequential(\n",
    "    [\n",
    "        layers.Dense(256, activation='relu',input_shape=(28*28,)),\n",
    "        layers.Dropout(0.2),\n",
    "        layers.Dense(128, activation='tanh'),\n",
    "        layers.Dropout(0.2),\n",
    "        layers.Dense(10, activation='softmax')\n",
    "    ]\n",
    ")\n",
    "\n",
    "# 编译模型 自适应矩估计\n",
    "model.compile(optimizer='adam',\n",
    "              loss='categorical_crossentropy',\n",
    "              metrics=['accuracy'])\n",
    "\n",
    "# 训练模型，验证集比例为0.2(帮助adam判断是否需要调整参数)，训练10轮\n",
    "history = model.fit(x_train, y_train,\n",
    "          batch_size=128,\n",
    "          epochs=10,\n",
    "          validation_split=0.2)\n",
    "\n",
    "test_loss,test_acc = model.evaluate(x_test, y_test)\n",
    "print(f'accuracy: {test_acc:.4f}')\n",
    "\n",
    "model.save('./models/mnist_model-tf.h5')  # 保存为HDF5格式\n",
    "print(\"模型已保存到 ./models/mnist_model-tf.h5\")\n",
    "\n",
    "#plt绘制\n",
    "plt.figure(figsize=(12, 4))\n",
    "\n",
    "plt.subplot(1, 2, 1)\n",
    "plt.plot(history.history['accuracy'], label='train accuracy')\n",
    "plt.plot(history.history['val_accuracy'], label='verify accuracy')\n",
    "plt.title('accuracy')\n",
    "plt.xlabel('Epoch')\n",
    "plt.ylabel('Accuracy')\n",
    "plt.legend()\n",
    "\n",
    "plt.subplot(1, 2, 2)\n",
    "plt.plot(history.history['loss'], label='train loss')\n",
    "plt.plot(history.history['val_loss'], label='val loss')\n",
    "plt.title('loss')\n",
    "plt.xlabel('Epoch')\n",
    "plt.ylabel('Loss')\n",
    "plt.legend()\n",
    "\n",
    "plt.show()\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "2025-03-14 18:57:55.832397: I metal_plugin/src/device/metal_device.cc:1154] Metal device set to: Apple M1\n",
      "2025-03-14 18:57:55.832421: I metal_plugin/src/device/metal_device.cc:296] systemMemory: 16.00 GB\n",
      "2025-03-14 18:57:55.832428: I metal_plugin/src/device/metal_device.cc:313] maxCacheSize: 5.33 GB\n",
      "2025-03-14 18:57:55.832455: I tensorflow/core/common_runtime/pluggable_device/pluggable_device_factory.cc:306] Could not identify NUMA node of platform GPU ID 0, defaulting to 0. Your kernel may not have been built with NUMA support.\n",
      "2025-03-14 18:57:55.832469: I tensorflow/core/common_runtime/pluggable_device/pluggable_device_factory.cc:272] Created TensorFlow device (/job:localhost/replica:0/task:0/device:GPU:0 with 0 MB memory) -> physical PluggableDevice (device: 0, name: METAL, pci bus id: <undefined>)\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 13944 characters]",
      "text/plain": [
       "<Figure size 640x480 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "预处理后的数据范围： 0.12156863 0.83137256\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "2025-03-14 18:57:56.529571: I tensorflow/core/grappler/optimizers/custom_graph_optimizer_registry.cc:117] Plugin optimizer for device_type GPU is enabled.\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "1/1 [==============================] - 1s 753ms/step\n",
      "预测概率分布： [[9.7845614e-01 1.6016074e-06 5.1186988e-03 4.2060162e-03 3.1019366e-07\n",
      "  1.1911159e-02 5.7799196e-05 6.7962850e-05 3.7460737e-05 1.4272679e-04]]\n",
      "预测结果：0\n"
     ]
    }
   ],
   "source": [
    "# 测试tf训练的手写数据识别\n",
    "import tensorflow as tf\n",
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "\n",
    "# 加载模型\n",
    "model = tf.keras.models.load_model('./models/mnist_model-tf.h5')\n",
    "\n",
    "# 加载并预处理图片\n",
    "image_path = './test/0.png'\n",
    "image = tf.keras.preprocessing.image.load_img(image_path, target_size=(28, 28), color_mode='grayscale')\n",
    "image = tf.keras.preprocessing.image.img_to_array(image)\n",
    "\n",
    "# 可视化预处理后的图片\n",
    "plt.imshow(image, cmap='gray')\n",
    "plt.title('Processed Image')\n",
    "plt.show()\n",
    "\n",
    "# 进一步处理\n",
    "image = image.reshape(-1, 28 * 28).astype('float32') / 255.0\n",
    "\n",
    "# 打印预处理后的数据\n",
    "print(\"预处理后的数据范围：\", np.min(image), np.max(image))\n",
    "\n",
    "# 进行预测\n",
    "predictions = model.predict(image)\n",
    "print(\"预测概率分布：\", predictions)\n",
    "\n",
    "# 获取预测结果\n",
    "predicted_class = tf.argmax(predictions, axis=1).numpy()[0]\n",
    "print(f\"预测结果：{predicted_class}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 仅全链接层手写识别 - torch\n",
    "\n",
    "## 模型定义\n",
    "- 输入层：28*28=784\n",
    "- 输出层：10\n",
    "- 隐藏层：784(input)->256(relu)(dropout-20%)->128(tanh)(dropout-20%)->10(softmax)\n",
    "\n",
    "## 数据处理\n",
    "- torchvision中提供了mnist数据集的下载和处理(转换为tensor并进行归一化)\n",
    "- 使用dataloader对数据进行分批处理\n",
    "\n",
    "## 定义模型\n",
    "\n",
    "- torch对模型的定义更加细化，需要继承nn.Module类\n",
    "- 定义模型的构造函数，定义模型的层\n",
    "- 定义模型的前向传播函数，定义模型的层之间的连接关系和数据的流动方向\n",
    "\n",
    "## 训练\n",
    "\n",
    "- 设置设备\n",
    "- 定义损失函数为交叉熵损失函数\n",
    "- 定义优化器为Adam优化器\n",
    "- 与tensorflow不同的是，torch需要手动实现训练过程\n",
    "- 并且，torch中的adam优化器不会根据训练集或者验证集对学习率进行调整，梯度的一阶矩（均值）和二阶矩（方差）​动态调整每个参数的学习率，如果想让其根据进行特定的数据集学习率调整，需要加入学习调度器的支持\n",
    "\n",
    "## 具体训练\n",
    "\n",
    "- 每一个batch中，需要将梯度清零\n",
    "- 每个batch中的image就是一次训练，需要将其传入指定设备\n",
    "- 每个image轮次都需要先试用优化器对每个参数的梯度清零，因为torch的梯度计算本质是对每个参数的梯度累加\n",
    "- 前向传播出结果\n",
    "- 使用loss函数计算损失，使用loss.backward自动梯度计算并存储（累加）在每个权重的.grad属性中\n",
    "- 最后使用优化器对每个参数的梯度进行更新\n",
    "\n",
    "## 保存\n",
    "\n",
    "- torch模型的保存需要手动指定保存的内容，如模型的参数、优化器的状态、损失函数的状态等，类似于键值对的形式存入模型文件\n",
    "\n",
    "## 预测\n",
    "\n",
    "- torch的模型加载使用torch.load函数可以直接加载，但是这个并非是模型实例，而是一个字典，需要手动指定加载的内容以加载出模型，然后再对模型进行参数的加载\n",
    "- 需要将模型设置为评估模式，否则会出现梯度计算的问题\n",
    "- 图片传入后可以使用torchvision的transforms进行处理（分辨率，灰度图，向量化，归一化）\n",
    "- 预测结果为tensor，需要使用torch.argmax函数获取最大值的索引，即为预测结果（注意预测时要在非自动梯度计算下，节省性能内存并防止意外更新）\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Epoch 1/10: 100%|██████████| 938/938 [00:10<00:00, 87.50batch/s, loss=1.6135, acc=86.53%] \n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 1/10 => Loss: 1.6118, Accuracy: 86.53%\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Epoch 2/10: 100%|██████████| 938/938 [00:10<00:00, 92.81batch/s, loss=1.5511, acc=92.52%] \n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 2/10 => Loss: 1.5395, Accuracy: 92.52%\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Epoch 3/10: 100%|██████████| 938/938 [00:09<00:00, 96.63batch/s, loss=1.5393, acc=93.52%] \n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 3/10 => Loss: 1.5278, Accuracy: 93.52%\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Epoch 4/10: 100%|██████████| 938/938 [00:09<00:00, 97.68batch/s, loss=1.5318, acc=94.25%] \n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 4/10 => Loss: 1.5204, Accuracy: 94.25%\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Epoch 5/10: 100%|██████████| 938/938 [00:09<00:00, 97.31batch/s, loss=1.5334, acc=94.56%] \n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 5/10 => Loss: 1.5171, Accuracy: 94.56%\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Epoch 6/10: 100%|██████████| 938/938 [00:09<00:00, 97.19batch/s, loss=1.5135, acc=95.05%] \n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 6/10 => Loss: 1.5119, Accuracy: 95.05%\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Epoch 7/10: 100%|██████████| 938/938 [00:09<00:00, 96.52batch/s, loss=1.5211, acc=95.24%] \n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 7/10 => Loss: 1.5097, Accuracy: 95.24%\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Epoch 8/10: 100%|██████████| 938/938 [00:09<00:00, 99.28batch/s, loss=1.5137, acc=95.30%] \n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 8/10 => Loss: 1.5089, Accuracy: 95.30%\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Epoch 9/10: 100%|██████████| 938/938 [00:09<00:00, 99.06batch/s, loss=1.5127, acc=95.55%] \n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 9/10 => Loss: 1.5063, Accuracy: 95.55%\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Epoch 10/10: 100%|██████████| 938/938 [00:09<00:00, 100.77batch/s, loss=1.5100, acc=95.68%]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 10/10 => Loss: 1.5052, Accuracy: 95.68%\n",
      "模型已保存到 ./models/mnist_model-torch.pth\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 73212 characters]",
      "text/plain": [
       "<Figure size 1200x500 with 2 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "import torch\n",
    "import torch.nn as nn\n",
    "import torch.optim as optim\n",
    "from tqdm import tqdm\n",
    "from torchvision import datasets,transforms\n",
    "import matplotlib.pyplot as plt\n",
    "\n",
    "# 处理数据\n",
    "transform = transforms.Compose([\n",
    "    transforms.ToTensor(),\n",
    "    transforms.Normalize(mean=[0.5],std=[0.5])\n",
    "])\n",
    "\n",
    "# 导入数据\n",
    "trainset = datasets.MNIST(root='./data',train=True,download=True,transform=transform)\n",
    "trainloader = torch.utils.data.DataLoader(trainset,batch_size=64,shuffle=True)\n",
    "\n",
    "trainset = datasets.MNIST(root='./data',train=True,download=True,transform=transform)\n",
    "testloader = torch.utils.data.DataLoader(trainset,batch_size=64,shuffle=True)\n",
    "\n",
    "# 定义模型\n",
    "class SimpleNet(nn.Module):\n",
    "    def __init__(self):\n",
    "        super(SimpleNet,self).__init__()\n",
    "        self.flatten = nn.Flatten()\n",
    "        self.fc1 = nn.Linear(28*28,256)\n",
    "        self.d1 = nn.Dropout(0.2)\n",
    "        self.fc2 = nn.Linear(256,128)\n",
    "        self.d2 = nn.Dropout(0.2)\n",
    "        self.fc3 = nn.Linear(128,10)\n",
    "        self.softmax = nn.Softmax(dim=1)\n",
    "\n",
    "    def forward(self,x):\n",
    "        x = self.flatten(x)\n",
    "        x = torch.relu(self.fc1(x))\n",
    "        x = self.d1(x)\n",
    "        x = torch.tanh(self.fc2(x))\n",
    "        x = self.d2(x)\n",
    "        x = self.fc3(x)\n",
    "        x = self.softmax(x)\n",
    "        return x    \n",
    "    \n",
    "model = SimpleNet()\n",
    "\n",
    "device = torch.device(\"mps\" if torch.backends.mps.is_available() else \"cpu\")\n",
    "\n",
    "model.to(device)\n",
    "\n",
    "# 定义损失函数和优化器\n",
    "criterion = nn.CrossEntropyLoss()\n",
    "optimizer = optim.Adam(model.parameters(),lr=0.001)\n",
    "\n",
    "# 初始化训练统计信息\n",
    "train_losses = []\n",
    "train_accuracies = []\n",
    "\n",
    "# 训练模型\n",
    "epochs = 10\n",
    "for epoch in range(epochs):\n",
    "    running_loss = 0.0\n",
    "    correct = 0\n",
    "    total = 0\n",
    "    \n",
    "    # 添加进度条\n",
    "    train_iter = tqdm(trainloader, desc=f'Epoch {epoch+1}/{epochs}', unit='batch')\n",
    "    \n",
    "    for images, labels in train_iter:\n",
    "        images, labels = images.to(device), labels.to(device)\n",
    "        optimizer.zero_grad()\n",
    "        \n",
    "        outputs = model(images)\n",
    "        loss = criterion(outputs, labels)\n",
    "        loss.backward()\n",
    "        optimizer.step()\n",
    "        \n",
    "        # 计算准确率\n",
    "        _, predicted = torch.max(outputs.data, 1)\n",
    "        total += labels.size(0)\n",
    "        correct += (predicted == labels).sum().item()\n",
    "        \n",
    "        # 更新进度条信息\n",
    "        running_loss += loss.item()\n",
    "        train_iter.set_postfix({\n",
    "            'loss': f\"{running_loss/(train_iter.n+1):.4f}\",\n",
    "            'acc': f\"{100*correct/total:.2f}%\"\n",
    "        })\n",
    "\n",
    "        # 保存训练统计信息\n",
    "        epoch_loss = running_loss / len(trainloader)\n",
    "        epoch_acc = 100 * correct / total\n",
    "        train_losses.append(epoch_loss)\n",
    "        train_accuracies.append(epoch_acc)\n",
    "\n",
    "# 在测试集上评估模型\n",
    "model.eval()  # 设置为评估模式\n",
    "test_loss = 0.0\n",
    "test_correct = 0\n",
    "test_total = 0\n",
    "\n",
    "with torch.no_grad():\n",
    "    for images, labels in testloader:\n",
    "        images, labels = images.to(device), labels.to(device)\n",
    "        outputs = model(images)\n",
    "        loss = criterion(outputs, labels)\n",
    "        \n",
    "        # 计算准确率\n",
    "        _, predicted = torch.max(outputs.data, 1)\n",
    "        test_total += labels.size(0)\n",
    "        test_correct += (predicted == labels).sum().item()\n",
    "        test_loss += loss.item()\n",
    "\n",
    "# 打印测试集评估结果\n",
    "test_loss /= len(testloader)\n",
    "test_acc = 100 * test_correct / test_total\n",
    "print(f\"测试集评估结果 => Loss: {test_loss:.4f}, Accuracy: {test_acc:.2f}%\")\n",
    "\n",
    "# 保存模型\n",
    "torch.save({\n",
    "    'model_state_dict': model.state_dict(),\n",
    "    'optimizer_state_dict': optimizer.state_dict(),\n",
    "    'model_definition': SimpleNet  # 保存模型类\n",
    "}, './models/mnist_model-torch.pth')\n",
    "print(\"模型已保存到 ./models/mnist_model-torch.pth\")\n",
    "\n",
    "# 训练结束后绘制曲线\n",
    "plt.figure(figsize=(12, 5))\n",
    "\n",
    "# 绘制损失曲线\n",
    "plt.subplot(1, 2, 1)\n",
    "plt.plot(train_losses, label='Train Loss')\n",
    "plt.xlabel('Epoch')\n",
    "plt.ylabel('Loss')\n",
    "plt.title('Training Loss')\n",
    "plt.legend()\n",
    "\n",
    "\n",
    "# 绘制准确率曲线\n",
    "plt.subplot(1, 2, 2)\n",
    "plt.plot(train_accuracies, label='Train Accuracy')\n",
    "plt.xlabel('Epoch')\n",
    "plt.ylabel('Accuracy')\n",
    "plt.title('Training Accuracy')\n",
    "plt.legend()\n",
    "\n",
    "plt.tight_layout()\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/var/folders/cf/pykndz2n2tl6ff6jl9bcbb480000gn/T/ipykernel_3293/2215973287.py:8: FutureWarning: You are using `torch.load` with `weights_only=False` (the current default value), which uses the default pickle module implicitly. It is possible to construct malicious pickle data which will execute arbitrary code during unpickling (See https://github.com/pytorch/pytorch/blob/main/SECURITY.md#untrusted-models for more details). In a future release, the default value for `weights_only` will be flipped to `True`. This limits the functions that could be executed during unpickling. Arbitrary objects will no longer be allowed to be loaded via this mode unless they are explicitly allowlisted by the user via `torch.serialization.add_safe_globals`. We recommend you start setting `weights_only=True` for any use case where you don't have full control of the loaded file. Please open an issue on GitHub for any issues related to this experimental feature.\n",
      "  checkpoint = torch.load('./models/mnist_model-torch.pth')\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 14848 characters]",
      "text/plain": [
       "<Figure size 640x480 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "预测概率分布： [[0.08533675 0.08533675 0.2319692  0.0853368  0.08533675 0.08533675\n",
      "  0.08533675 0.08533675 0.08533676 0.08533675]]\n",
      "预测结果：2\n"
     ]
    }
   ],
   "source": [
    "# 测试torch训练的手写数据识别\n",
    "import torch\n",
    "from PIL import Image\n",
    "import torchvision.transforms as transforms\n",
    "import matplotlib.pyplot as plt\n",
    "\n",
    "# 加载模型\n",
    "checkpoint = torch.load('./models/mnist_model-torch.pth')\n",
    "model = checkpoint['model_definition']()\n",
    "model.load_state_dict(checkpoint['model_state_dict'])\n",
    "model.eval()  # 设置为评估模式\n",
    "\n",
    "# 图片预处理\n",
    "transform = transforms.Compose([\n",
    "    transforms.Resize((28, 28)),\n",
    "    transforms.Grayscale(),\n",
    "    transforms.ToTensor(),\n",
    "    transforms.Normalize((0.5,), (0.5,))\n",
    "])\n",
    "\n",
    "# 加载并预处理图片\n",
    "image_path = './test/2.png'\n",
    "image = Image.open(image_path)\n",
    "processed_image = transform(image)\n",
    "\n",
    "# 可视化预处理后的图片\n",
    "plt.imshow(processed_image.squeeze(), cmap='gray')\n",
    "plt.title('Processed Image')\n",
    "plt.show()\n",
    "\n",
    "# 进行预测\n",
    "with torch.no_grad():\n",
    "    output = model(processed_image.unsqueeze(0))\n",
    "    probabilities = torch.nn.functional.softmax(output, dim=1)\n",
    "    predicted_class = torch.argmax(probabilities, dim=1).item()\n",
    "\n",
    "print(\"预测概率分布：\", probabilities.numpy())\n",
    "print(f\"预测结果：{predicted_class}\")\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/var/folders/cf/pykndz2n2tl6ff6jl9bcbb480000gn/T/ipykernel_3293/3759706733.py:8: FutureWarning: You are using `torch.load` with `weights_only=False` (the current default value), which uses the default pickle module implicitly. It is possible to construct malicious pickle data which will execute arbitrary code during unpickling (See https://github.com/pytorch/pytorch/blob/main/SECURITY.md#untrusted-models for more details). In a future release, the default value for `weights_only` will be flipped to `True`. This limits the functions that could be executed during unpickling. Arbitrary objects will no longer be allowed to be loaded via this mode unless they are explicitly allowlisted by the user via `torch.serialization.add_safe_globals`. We recommend you start setting `weights_only=True` for any use case where you don't have full control of the loaded file. Please open an issue on GitHub for any issues related to this experimental feature.\n",
      "  checkpoint = torch.load('./models/mnist_model-torch.pth')\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 19304 characters]",
      "text/plain": [
       "<Figure size 1200x400 with 3 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "图片 ./test/0.png 的预测结果：0\n",
      "概率分布：[0.23196927 0.08533674 0.08533675 0.08533674 0.08533674 0.08533674\n",
      " 0.08533674 0.08533675 0.08533674 0.08533674]\n",
      "----------------------------------------\n",
      "图片 ./test/2.png 的预测结果：2\n",
      "概率分布：[0.08533675 0.08533675 0.2319692  0.0853368  0.08533675 0.08533675\n",
      " 0.08533675 0.08533675 0.08533676 0.08533675]\n",
      "----------------------------------------\n",
      "图片 ./test/3.png 的预测结果：3\n",
      "概率分布：[0.08533674 0.08533674 0.08533674 0.23196931 0.08533674 0.08533674\n",
      " 0.08533674 0.08533674 0.08533674 0.08533674]\n",
      "----------------------------------------\n"
     ]
    }
   ],
   "source": [
    "# 多图片同时预测\n",
    "import torch\n",
    "from PIL import Image\n",
    "import torchvision.transforms as transforms\n",
    "import matplotlib.pyplot as plt\n",
    "\n",
    "# 加载模型\n",
    "checkpoint = torch.load('./models/mnist_model-torch.pth')\n",
    "model = checkpoint['model_definition']()\n",
    "model.load_state_dict(checkpoint['model_state_dict'])\n",
    "model.eval()  # 设置为评估模式\n",
    "\n",
    "# 图片预处理\n",
    "transform = transforms.Compose([\n",
    "    transforms.Resize((28, 28)),\n",
    "    transforms.Grayscale(),\n",
    "    transforms.ToTensor(),\n",
    "    transforms.Normalize((0.5,), (0.5,))\n",
    "])\n",
    "\n",
    "# 加载并预处理多个图片\n",
    "image_paths = ['./test/0.png', './test/2.png', './test/3.png']  # 添加更多图片路径\n",
    "images = [Image.open(path) for path in image_paths]\n",
    "processed_images = torch.stack([transform(img) for img in images])  # 将多个图片堆叠成一个批次\n",
    "\n",
    "# 可视化预处理后的图片\n",
    "fig, axes = plt.subplots(1, len(images), figsize=(12, 4))\n",
    "for i, img in enumerate(processed_images):\n",
    "    axes[i].imshow(img.squeeze(), cmap='gray')\n",
    "    axes[i].set_title(f'Image {i+1}')\n",
    "plt.show()\n",
    "\n",
    "# 进行预测\n",
    "with torch.no_grad():\n",
    "    outputs = model(processed_images)  # 直接传入批次数据\n",
    "    probabilities = torch.nn.functional.softmax(outputs, dim=1)\n",
    "    predicted_classes = torch.argmax(probabilities, dim=1).numpy()\n",
    "\n",
    "# 打印预测结果\n",
    "for i, (path, pred, prob) in enumerate(zip(image_paths, predicted_classes, probabilities)):\n",
    "    print(f\"图片 {path} 的预测结果：{pred}\")\n",
    "    print(f\"概率分布：{prob.numpy()}\")\n",
    "    print(\"-\" * 40)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "ail",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.16"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
