- 2024-05-26
-
回复了主题帖:
#AI挑战营终点站#RV1106手写数字识别部署
luyism 发表于 2024-5-26 13:32
将图像的数据复制到张量内存空间的时候是要求16字节对齐的,28字节一行的图像数据直接整个拷贝到里面不满足 ...
好的我修改一下,感谢感谢
-
回复了主题帖:
【AI挑战营终点站】应用落地:部署手写数字识别应用到幸狐RV1106开发板
完成打卡#AI挑战营终点站#RV1106手写数字识别部署 https://bbs.eeworld.com.cn/thread-1282914-1-1.html
- 2024-05-25
-
发表了主题帖:
#AI挑战营终点站#RV1106手写数字识别部署
本帖最后由 天空中 于 2024-5-25 18:09 编辑
# RV1106数字识别
## 简述
利用 RKMPI 库实现摄像头图像捕获、预处理、硬件编码,结合 opencv-mobile 进行图像处理, rtsp 推流,使用 VLC 软件拉取并观察图像。
### 实现效果
* RKNN数字识别并标注
## 操作流程
RKMPI获取图像 -> 使用 opencv-mobile 进行捕获 ->opencv-mobile 进行图像预处理框选数字 -> RKNN推理->结果反回
## 涉及函数库简述
### opencv-mobile
在嵌入式系统中opencv中很多模块无法使用如:opencv_gapi,opencv_videoio等模块,所以需要精简openncv。
opencv-mobile,一款体积仅有官方 1/10 的精简 OpenCV 库,并已经适配rv1106已支持 luckfox-pico MIPI CSI 摄像头和 rk-aiq/rga 硬件加速
[github库地址](https://github.com/nihui/opencv-mobile)
### RKNN模型转换并部署
RKNN-Toolkit2 工具在 PC 平台上提供 C 或 Python 接口,简化模型的部署和运行。用户可以通过该工具轻松完成以下功能:模型转换、量化、推理、性能和内存评估、量化精度分析以及模型加密。RKNN-Toolkit2的环境及安装在以前的文章中有详细讲述此文章不在赘述
## 代码讲解
感谢 [knv](https://home.eeworld.com.cn/space-uid-1362174.html) 的代码分享,提供了思路和优化手段。
本代码参考文章[RKMPI 实例使用指南 | LUCKFOX WIKI](https://wiki.luckfox.com/zh/Luckfox-Pico/RKMPI-example#6-opencv-mobile-标注推流帧数实例)
### 将RKMPI获取的图像转化为opencv Mat型
```c++
void *data = RK_MPI_MB_Handle2VirAddr(stVpssFrame.stVFrame.pMbBlk); //获取MPI数据帧对应内存缓冲块虚拟地址
cv::Mat frame(height,width,CV_8UC3, data);//转化为opencv Mat型为后序图像预处理做准备
```
### 图像预处理
```c++
cv::Mat frame64;
cv::cvtColor(frame, frame64, cv::COLOR_BGR2GRAY); //将彩色图转为灰度图
//图形二值化,此处没有使用固定阈值而是使用大津法做动态阈值处理增加了在不同光源下识别数字的稳定性
//使用THRESH_BINARY_INV模式将大于阈值的像素设为0小于阈值设为255这样就不用在进行像素反转的操作
cv::threshold(frame64, frame64, 100, 255, cv::THRESH_BINARY_INV | cv::THRESH_OTSU);
```
#### 寻找数字的轮廓并框选
```c++
// 查找轮廓,声明变量来存储轮廓
std::vectorcontours;
std::vector hierarchy;
cv::findContours(frame64, contours, hierarchy, cv::RETR_EXTERNAL, cv::CHAIN_APPROX_NONE, cv::Point());
//声明变量来存储框选出来的矩形
cv::Rect rect;
int out_num = 10;
for (int i = 0; i < contours.size(); i++)
{
rect = cv::boundingRect(contours); //在寻找的轮廓中拟合矩形
//通过矩形的长度宽度和长宽比筛出错误的矩形
if (rect.width < 100 && rect.width > 10 && (rect.height/rect.width)>1 && rect.height virt_addr, roi.data, 28*28*1);
//运行模型
inference_NUM_model(&rknn_app_ctx, out_num);
}
}
```
本代码没有对图像进行滤波操作,因为经过测试滤波操作在光线稳定的情况下并不会影响数字的识别率并且滤波操作会遍历整张图像降低帧率所以舍去滤波操作。
## 模型初始化部分和模型推理部分使用RKNN零拷贝 API
注:rv1106/rv1103仅支持零拷贝 API
RKNN零拷贝API是Rockchip神经网络处理库(RKNN, Rockchip Neural Network)提供的一组功能,旨在优化数据传输效率,减少内存复制操作,从而提高基于Rockchip平台的深度学习应用的性能。在传统的内存操作中,数据在不同内存区域或不同进程间传递时,经常需要进行数据的复制,这一过程会消耗时间和系统资源。而零拷贝技术则尽量避免这些不必要的复制操作。
### 模型初始化
```c++
//rknn上下文结构体变量
typedef struct {
// RKNN运行上下文,由rknn_init函数初始化后获得,是执行模型推理的核心句柄。
rknn_context rknn_ctx;
// 指向最大内存块的指针,可根据需要动态调整以适应不同大小的输入或中间结果。
rknn_tensor_mem* max_mem;
// 网络内存指针,通常用于存储模型运行过程中的临时数据或中间结果。
rknn_tensor_mem* net_mem;
// 输入输出数量结构体,记录模型的输入和输出张量数量。
rknn_input_output_num io_num;
// 指向输入张量属性数组的指针,存储每个输入张量的详细属性信息(如维度、数据类型等)。
rknn_tensor_attr* input_attrs;
// 指向输出张量属性数组的指针,存储每个输出张量的详细属性信息。
rknn_tensor_attr* output_attrs;
// 输入内存块指针数组,至少包含一个元素,用于存放模型推理的输入数据。
rknn_tensor_mem* input_mems[1];
// 输出内存块指针数组,这里假设模型有3个输出,用于接收模型推理的结果数据。
rknn_tensor_mem* output_mems[3];
// 模型输入通道数,对于图像处理任务,这通常指的是色彩通道数(如RGB图像的通道数为3)。
int model_channel;
// 模型输入宽度,单位通常是像素。
int model_width;
// 模型输入高度,单位通常是像素。
int model_height;
// 表示模型是否进行了量化处理。量化模型通常用于加速推理并减少内存占用。
bool is_quant;
} rknn_app_context_t;
```
```c++
rknn_app_context_t rknn_app_ctx; // rknn结构体变量
int ret;
const char *model_path = "./model/111.rknn"; //模型地址
memset(&rknn_app_ctx, 0, sizeof(rknn_app_context_t)); //初始化结构体
init_retinaface_model(model_path, &rknn_app_ctx); //模型初始化函数
printf("init rknn model success!\n");
```
```c++
/**
* 初始化RKNN模型的函数
*
* @param model_path 模型文件的路径
* @param app_ctx 应用上下文结构体指针,用于存储模型上下文及输入输出相关信息
* @return 成功返回0,失败返回-1
*/
int init_retinaface_model(const char *model_path, rknn_app_context_t *app_ctx)
{
int ret;
rknn_context ctx = 0; // RKNN模型运行上下文
// 使用rknn_init函数初始化模型,传入模型路径
ret = rknn_init(&ctx, (char *)model_path, 0, 0, NULL);
if (ret < 0)
{
printf("rknn_init fail! ret=%d\n", ret);
return -1;
}
// 查询模型的输入输出数量
rknn_input_output_num io_num;
ret = rknn_query(ctx, RKNN_QUERY_IN_OUT_NUM, &io_num, sizeof(io_num));
if (ret != RKNN_SUCC)
{
printf("rknn_query fail! ret=%d\n", ret);
return -1;
}
printf("model input num: %d, output num: %d\n", io_num.n_input, io_num.n_output);
// 分别获取并打印输入和输出张量的详细信息
rknn_tensor_attr input_attrs[io_num.n_input];
memset(input_attrs, 0, sizeof(input_attrs));
for (int i = 0; i < io_num.n_input; i++) {
input_attrs.index = i;
ret = rknn_query(ctx, RKNN_QUERY_NATIVE_INPUT_ATTR, &(input_attrs), sizeof(rknn_tensor_attr));
if (ret != RKNN_SUCC) return -1;
dump_tensor_attr(&(input_attrs)); // 打印输入张量属性
}
rknn_tensor_attr output_attrs[io_num.n_output];
memset(output_attrs, 0, sizeof(output_attrs));
for (int i = 0; i < io_num.n_output; i++) {
output_attrs.index = i;
ret = rknn_query(ctx, RKNN_QUERY_NATIVE_NHWC_OUTPUT_ATTR, &(output_attrs), sizeof(rknn_tensor_attr));
if (ret != RKNN_SUCC) return -1;
dump_tensor_attr(&(output_attrs)); // 打印输出张量属性
}
// 设置输入张量属性,并为输入分配内存
input_attrs[0].type = RKNN_TENSOR_UINT8;
input_attrs[0].fmt = RKNN_TENSOR_NHWC;
printf("input_attrs[0].size_with_stride=%d\n", input_attrs[0].size_with_stride);
app_ctx->input_mems[0] = rknn_create_mem(ctx, input_attrs[0].size_with_stride);
ret = rknn_set_io_mem(ctx, app_ctx->input_mems[0], &input_attrs[0]);
// 为输出张量分配内存并设置内存属性
for (uint32_t i = 0; i < io_num.n_output; ++i) {
app_ctx->output_mems = rknn_create_mem(ctx, output_attrs.size_with_stride);
printf("output mem [%d] = %d \n",i ,output_attrs.size_with_stride);
ret = rknn_set_io_mem(ctx, app_ctx->output_mems, &output_attrs);
if (ret < 0) return -1;
}
// 将模型上下文和其他相关信息保存到app_ctx中
app_ctx->rknn_ctx = ctx;
app_ctx->is_quant = (output_attrs[0].qnt_type == RKNN_TENSOR_QNT_AFFINE_ASYMMETRIC);
app_ctx->io_num = io_num;
// 分配内存并复制输入输出属性到app_ctx中
app_ctx->input_attrs = (rknn_tensor_attr *)malloc(io_num.n_input * sizeof(rknn_tensor_attr));
memcpy(app_ctx->input_attrs, input_attrs, io_num.n_input * sizeof(rknn_tensor_attr));
app_ctx->output_attrs = (rknn_tensor_attr *)malloc(io_num.n_output * sizeof(rknn_tensor_attr));
memcpy(app_ctx->output_attrs, output_attrs, io_num.n_output * sizeof(rknn_tensor_attr));
// 根据输入格式确定模型的尺寸信息
if (input_attrs[0].fmt == RKNN_TENSOR_NCHW)
{
printf("model is NCHW input fmt\n");
app_ctx->model_channel = input_attrs[0].dims[1];
app_ctx->model_height = input_attrs[0].dims[2];
app_ctx->model_width = input_attrs[0].dims[3];
} else
{
printf("model is NHWC input fmt\n");
app_ctx->model_height = input_attrs[0].dims[1];
app_ctx->model_width = input_attrs[0].dims[2];
app_ctx->model_channel = input_attrs[0].dims[3];
}
printf("model input height=%d, width=%d, channel=%d\n",
app_ctx->model_height, app_ctx->model_width, app_ctx->model_channel);
return 0; // 成功返回
}
```
### 模型推理函数
```c++
/**
* 执行模型推理并处理输出结果的函数
*
* @param app_ctx 指向RKNN应用上下文的指针,包含了模型运行所需的全部资源和参数
* @param out_num 引用参数,推理后最可能的类别序号将被保存于此
* @return 成功返回0,失败返回-1
*/
int inference_NUM_model(rknn_app_context_t *app_ctx, int &out_num)
{
int ret = 0;
float probability = 0;
// 使用rknn_run函数执行模型推理
ret = rknn_run(app_ctx->rknn_ctx, NULL);
if (ret < 0)
{
printf("rknn run error %d\n", ret);
return -1;
}
// 获取模型输出数据的指针
int8_t* output_data = (int8_t*)(app_ctx->output_mems[0]->virt_addr);
// 创建一个数组用于存放反量化后的输出值
float dequantized_values[10];
// 对输出数据进行反量化处理
for (int i = 0; i < 10; ++i)
{
dequantized_values = (output_data - (float)(-4)) * 0.122913;
}
// 应用softmax函数使输出值转换为概率
applySoftmax(dequantized_values, 10);
for (int i = 0; i < 10; ++i)
{
printf("Class %d: Probability = %.5f\n", i, dequantized_values);
// 如果当前类别的概率大于0.5,则检查是否大于已找到的最大概率
if(dequantized_values > 0.5)
{
// 更新最大概率和对应的类别序号
if(probability < dequantized_values)
{
probability = dequantized_values;
out_num = i;
}
}
}
return ret;
}
```
### 输出值转换为概率
```c++
void applySoftmax(float* data, int size)
{
// 寻找数组中的最大值,用于数值稳定化softmax计算
float max_val = data[0];
for (int i = 1; i < size; ++i) {
if (data > max_val) {
max_val = data;
}
}
// 计算每个元素exp(xi - max_val)的值并累加,用于归一化
float sum_exp = 0;
for (int i = 0; i < size; ++i) {
data = exp(data - max_val); // 对每个值减去最大值后进行指数运算,防止因指数运算导致的下溢出
sum_exp += data; // 累加所有指数值,用于后续的归一化
}
// 归一化步骤,将所有exp值转换为概率分布(所有概率之和为1)
for (int i = 0; i < size; ++i) {
data /= sum_exp; // 每个元素除以总和,得到概率值
}
}
```
## 代码运行效果
[localvideo]68fd412d539aec90db007434d4d9208e[/localvideo]
## 结果分析
本模型只对数字6,0,7这种特征较为明显的数字有较好的识别率,分析应该是神经网络的框架太简单,训练时只使用了三层卷积(具体训练代码看前面的文章)。可以将神经网络框架换为MobileNetv3这种成熟结构(仅为猜想并未实践)。
## 代码分享
链接:https://pan.baidu.com/s/1aSX1x_t9iWn5hv4Raj2_ow?pwd=ihrk
提取码:ihrk
-
加入了学习《直播回放: 瑞萨新一代视觉 AI MPU 处理器 RZ/V2H: 高算力、低功耗、实时控制》,观看 瑞萨新一代视觉 AI MPU 处理器 RZ/V2H
- 2024-05-08
-
回复了主题帖:
入围名单公布:嵌入式工程师AI挑战营(初阶),获RV1106 Linux 板+摄像头的名单
个人信息已确认,领取板卡,可继续完成&分享挑战营第二站和第三站任务。
-
回复了主题帖:
入围名单公布:嵌入式工程师AI挑战营(初阶),获RV1106 Linux 板+摄像头的名单
个人信息已确认,领取板卡,可继续完成&分享挑战营第二站和第三站任务。
- 2024-04-30
-
回复了主题帖:
【AI挑战营第二站】算法工程化部署打包成SDK
1.ONNX是一种开放的神经网络交换格式,它的主要目标是提供一种通用的模型表示,使得深度学习模型可以在不同的框架之间无缝转换和使用。
RKNN是由瑞芯微电子公司推出的一款深度学习模型各阶段流程一体化开发和运行框架。
2.#AI挑战营第二站#Ubuntu下ONNX转RKNN - 嵌入式系统 - 电子工程世界-论坛 (eeworld.com.cn)
-
发表了主题帖:
#AI挑战营第二站#Ubuntu下ONNX转RKNN
# 使用rknn-Toolkit2完成模型转换及其环境部署
## rknn-Toolkit2介绍
RKNN-Toolkit2 是为用户提供在计算机上进行模型转换、推理和性能评估的开发套件
用户可以通过该工具轻松完成以下功能:模型转换、量化、推理、性能和内存评估、量化精度分析以及模型加密。
## 环境部署
本帖使用Conda 创建 python 虚拟环境完成rknn-Toolkit2的安装
### 安装Conda
```
#下载安装包
wget https://mirrors.tuna.tsinghua.edu.cn/anaconda/miniconda/Miniconda3-4.6.14-Linux-x86_64.sh
#设置安装包权限并安装
chmod 777 Miniconda3-4.6.14-Linux-x86_64.sh
bash Miniconda3-4.6.14-Linux-x86_64.sh
```
### 下载rknn-Toolkit2安装包,并在Conda环境中安装
```
#下载rknn-Toolkit2安装包
git clone https://github.com/rockchip-linux/rknn-toolkit2
```
#### 创建 RKNN-Toolkit2 Conda 环境
```
#创建环境,设置python版本为3.8
conda create -n RKNN-Toolkit2 python=3.8
#进入环境
conda activate RKNN-Toolkit2
```
创建环境并成功进入后终端变化
#### 安装rknn-Toolkit2依赖
```
#在Conda环境中进入文件夹
cd rknn-toolkit2
#安装依赖
pip install tf-estimator-nightly==2.8.0.dev2021122109
pip install -r rknn-toolkit2/packages/requirements_cp38-1.6.0.txt -i https://pypi.mirrors.ustc.edu.cn/simple/
#安装 RKNN-Toolkit2
pip install rknn-toolkit2/packages/rknn_toolkit2-1.6.0+81f21f4d-cp38-cp38-linux_x86_64.whl
```
#### 验证是否成功
输入图片中的命令没有报错则说明成功
## 模型转换代码
```
import cv2
import numpy as np
from rknn.api import RKNN
import os
if __name__ == '__main__':
# 模型部署平台
platform = 'rv1106'
#训练模拟时输入图片大小
Width = 28
Height = 28
# 此处改为自己的模型地址
MODEL_PATH = '/home/tiankong/文档/RKNN_NUM/NUM_ONNX.onnx'
# 导出模型地址
RKNN_MODEL_PATH = '/home/tiankong/文档/RKNN_NUM/RKNN_Num.rknn'
# 创建RKNN对象并在屏幕打印详细的日志信息
rknn = RKNN(verbose=True)
# 模型配置
# mean_values: 输入图像像素均值
# std_values: 输入图像像素标准差
# target_platform: 目标部署平台
# 本模型训练时输入图象为单通道
rknn.config(mean_values=[0], std_values=[255], target_platform=platform)
# 模型加载
print('--> Loading model')
ret = rknn.load_onnx(MODEL_PATH)
if ret != 0:
print('load model failed!')
exit(ret)
print('done')
# 构建 RKNN 模型
print('--> Building model')
#do_quantization:是否对模型进行量化。默认值为 True
ret = rknn.build(do_quantization=True, dataset="./image_data.txt")
if ret != 0:
print('build model failed.')
exit(ret)
print('done')
# 导出模型
ret = rknn.export_rknn(RKNN_MODEL_PATH)
#释放RKNN模型
rknn.release()
```
## RKNN模型
- 2024-04-17
-
回复了主题帖:
免费申请:幸狐 RV1106 Linux 开发板(带摄像头),助力AI挑战营应用落地
#AI挑战营第一站#PyTorch+MNIST数据集手写数字识别 https://bbs.eeworld.com.cn/thread-1278227-1-1.html
预期应用:实现人脸识别功能,并跟踪人脸
-
回复了主题帖:
【AI挑战营第一站】模型训练:在PC上完成手写数字模型训练,免费申请RV1106开发板
1.模型训练的本质是通过给定的数据和标签,不断调整模型的参数,使模型能够从数据中学习到特征和规律。在训练过程中,模型通过损失函数衡量预测值与真实标签之间的差异,并使用优化算法不断优化模型参数,使损失函数达到最小值。训练的最终结果是得到了一个经过优化的模型,能够对未见过的数据进行准确的预测。
2.PyTorch是一个基于Python的深度学习框架,它提供了灵活的张量计算和动态神经网络构建的功能。PyTorch支持Windows、Linux和macOS,支持CPU、GPU和TPU等运算平台
3.https://bbs.eeworld.com.cn/thread-1278227-1-1.html
-
发表了主题帖:
#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, cmap='gray', interpolation='none')
plt.title("Labels: {}".format(train_dataset.train_labels))
plt.xticks([])
plt.yticks([])
plt.show()
```
## 定义模型
```
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()
```
## 训练结果展示
## 模型转换
```
# 创建模型实例并加载预训练的模型参数
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)
```
## 训练模型分享