STM32部署神经网络-pytorch

训练模型

在STM32微控制器上部署神经网络时，模型大小是一个重要的考虑因素，因为微控制器的存储和计算资源有限。以下是一些减小模型大小和优化模型以适应STM32微控制器的策略：

1. 模型架构选择：
   选择轻量级的神经网络架构，如MobileNet、SqueezeNet或Tiny YOLO，这些架构专为资源受限的设备设计。

本次选择轻量级网络MobleNetV1，MobileNet 是一个为移动和嵌入式视觉应用设计的卷积神经网络（CNN）架构。它由 Google 开发，旨在在保持高准确率的同时减少计算资源的消耗。MobileNet 特别适合在计算能力有限的设备上运行，比如智能手机和嵌入式系统。

修建模型

2. 参数剪枝：
   在训练过程中或训练后，剪除不重要的权重或神经元，减少模型的参数数量。
   非结构化剪枝：通过正则化技术（如L1正则化）或后训练剪枝策略来实现。
   结构化剪枝：通过分析模型的通道或层的重要性，并移除不重要的部分来实现。
   增量剪枝：逐步增加剪枝比例，逐步优化模型。
   动态剪枝：在模型推理时动态地移除或保留网络的部分。参考的文章：
   http://t.csdnimg.cn/AdPYh
   http://t.csdnimg.cn/yhkSa
   本文选用结构化剪枝，因为它可以显著减少计算量和内存占用。（非结构化剪枝产生的稀疏性可能需要特殊的硬件支持（如支持稀疏矩阵运算的GPU）来充分利用其优势。）

import torch.nn.utils.prune as prune

DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")# 训练模型
# 设置超参数
learning_rate = 0.01
num_epochs = 2
best_acc = 0
model = MobileNetV1(num_classes = 3)
# Define the pruning parameters
pruning_rate = 0.5  # Percentage of weights to prune

# Specify the pruning method (e.g., L1Unstructured, RandomUnstructured, etc.)
pruning_method = prune.L1Unstructured

# Create a pruning instance for each layer you want to prune
pruning_instances = [
#     (model.conv_dw1[0], "weight", pruning_method, pruning_rate),
    (model.conv_dw2[0], "weight", pruning_method, pruning_rate),
    # Add other layers to prune
    (model.conv_dw3[0], "weight", pruning_method, pruning_rate),
    (model.conv_dw4[0], "weight", pruning_method, pruning_rate),
    (model.conv_dw5[0], "weight", pruning_method, pruning_rate),
#     (model.conv_dw6[0], "weight", pruning_method, pruning_rate),
#     (model.conv_dw7[0], "weight", pruning_method, pruning_rate),
#     (model.conv_dw8[0], "weight", pruning_method, pruning_rate)
]

# Apply pruning to the model
for module, name, method, rate in pruning_instances:
    method.apply(module, name=name, amount=rate)

# 删除修剪的重参数化缓冲区
# prune.remove(model.conv_dw1[0], 'weight')
prune.remove(model.conv_dw2[0], 'weight')
prune.remove(model.conv_dw3[0], 'weight')
prune.remove(model.conv_dw4[0], 'weight')
prune.remove(model.conv_dw5[0], 'weight')

# 删除其他修剪的层的重参数化缓冲区
# model.fc = nn.Linear(16384, num_classes)  # Modify the input dimension

# Test the pruned model with sample input ，
inputs = torch.randn(4, 3, 64, 64)  # Corrected input dimensions
output = model(inputs)
# prune_weights(model, prune_ratio=0.5)
# prune_model(model, prune_ratio=0.5)
# # 移除被剪枝的权重
# parameters_to_prune = [(module, 'weight') for module in model.modules()]
# prune.remove(parameters_to_prune)
model.to(DEVICE)
# 定义损失函数和优化器
criterion = nn.CrossEntropyLoss()
optimizer = optim.SGD(model.parameters(), lr=learning_rate)
save_path = '/kaggle/working/model.pth

使用剪枝可以缩小部署模型占用stm32的flash，但是这时候，flash一般还是很大，接下来对模型的通道数进行减小

减少通道数

在卷积层中减少通道数可以显著减小模型的大小。你可以尝试减少每个卷积层的输出通道数，特别是在较深的层中。你可以通过修改 _conv_st 和 _conv_dw 方法中的 out_channels 参数来实现。
MobileNetV1源码（torch版）：http://t.csdnimg.cn/VuykL
经过减少后的代码（可根据自己的需求减少通道数，需要注意全连接层及池化层的kernel_size要符合自己的数据的输入输出）
调整 kernel_size：kernel_size 以确保它不大于输入张量的大小。
调整 stride： 以确保池化窗口可以覆盖整个输入张量。
检查模型结构：
确保模型的卷积层和池化层之间的尺寸匹配。

import torch
import torch.nn as nn
import torch.utils.checkpoint as checkpoint

class MobileNetV1(nn.Module):
    def __init__(self, num_classes):
        super(MobileNetV1, self).__init__()
        self.conv1 = self._conv_st(3, 4, 2)
        self.conv_dw2 = self._conv_dw(4, 8, 2)
        self.conv_dw3 = self._conv_dw(8, 16, 2)
        self.conv_dw4 = self._conv_dw(16, 32, 2)
        self.conv_dw5 = self._conv_dw(32, 64, 2)
        self.avgpool = nn.AvgPool2d(kernel_size=2, stride=2)
        self.fc = nn.Linear(64, num_classes)

    def forward(self, x):
        x = self.conv1(x)
        x = checkpoint.checkpoint(self.conv_dw2, x)
        x = checkpoint.checkpoint(self.conv_dw3, x)
        x = checkpoint.checkpoint(self.conv_dw4, x)
        x = checkpoint.checkpoint(self.conv_dw5, x)
        x = self.avgpool(x)
        x = x.view(x.size(0), -1)
        x = self.fc(x)
        return x

    def _conv_st(self, in_channels, out_channels, stride):
        return nn.Sequential(
            nn.Conv2d(in_channels, out_channels, kernel_size=3, stride=stride, padding=1, bias=False),
#             nn.BatchNorm2d(out_channels),
            nn.ReLU(inplace=True),
        )

    def _conv_dw(self, in_channels, out_channels, stride):
        return nn.Sequential(
            nn.Conv2d(in_channels, in_channels, kernel_size=3, stride=stride, padding=1, groups=in_channels, bias=False),
#             nn.BatchNorm2d(in_channels),
            nn.ReLU(inplace=True),
            nn.Conv2d(in_channels, out_channels, kernel_size=1, stride=1, padding=0, bias=False),
#             nn.BatchNorm2d(out_channels),
            nn.ReLU(inplace=True),
        )

此时使用stm32cubemx分析模型可能出现的结果：
Total Ram: 2166192 B (2.07 MiB)
Activations: 2160000 B (2.06 MiB)
Library: 6192 B (6.05 KiB)
Input: 1080000 B (1.03 MiB included in Activations)
模型的总 RAM 占用为 2,166,192 字节（大约 2.07 MiB），其中大部分内存被激活（Activations）占用，即 2,160,000 字节（大约 2.06 MiB）。输入数据（Input）占用 1.03 MiB，并且这部分内存被包含在激活内存中。

以下是一些优化内存使用的建议：

1. 优化激活内存：

   • 检查模型的激活函数和中间层的输出，看是否可以减少激活的大小。
   • 考虑使用更小的批量大小（batch size）来减少每次前向传播时的内存占用。

2. 减少输入数据大小：

   • 如果可能，减少输入数据的维度，例如降低图像的分辨率或减少序列的长度。

3. 模型简化：

   • 使用更小的模型或减少模型的深度和宽度。

4. 参数共享：

   • 在模型的不同部分之间共享权重，减少模型的参数数量。

5. 模型剪枝：

   • 移除模型中不重要的权重或神经元，减少模型的复杂性和内存占用。

6. 使用轻量级模型：

   • 考虑使用为资源受限环境设计的轻量级模型架构，如 MobileNets。

7. 内存管理：

   • 确保在训练循环中及时释放不再需要的变量和内存。

8. 量化模型：

   • 将模型的权重和激活从32位浮点数转换为更低位数（如8位或16位）。

9. 使用外部存储：

   • 对于非常大的模型或数据集，考虑使用外部存储（如硬盘或SSD）来存储模型或数据。

10. 硬件升级：

    • 如果硬件资源是限制因素，考虑升级硬件以提供更多的 RAM。

11. 代码优化：

    • 优化代码逻辑，减少不必要的内存分配和复制。

12. 监控内存使用：

    • 使用内存分析工具来监控内存使用情况，及时发现和解决内存泄漏问题。

通过这些方法，可以有效地减少模型的内存占用，提高其在资源受限环境中的可用性。
本文使用的方法是减小输入

减少输入数据大小

# 定义图像转换操作
transform = transforms.Compose([
    transforms.Resize((64, 64)),  # 降低图像分辨率
    transforms.ToTensor(),
    transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
])

模型转换与部署

参考资料

参考资料：https://blog.csdn.net/qq_48691686/article/details/131438689
http://t.csdnimg.cn/TThXv
https://blog.csdn.net/best_xiaolong/article/details/122053271
https://blog.csdn.net/xddwg521125/article/details/135358192