深度学习交通标志识别模型轻量化部署实战：从论文到生产环境

1次阅读

共计 4135 个字符，预计需要花费 11 分钟才能阅读完成。

交通标志识别（Traffic Sign Recognition, TSR）是自动驾驶和辅助驾驶系统的核心功能之一。准确识别道路上的交通标志可以帮助车辆做出正确的驾驶决策，提高行车安全性。然而，传统的深度学习模型如 ResNet、VGG 等通常参数量大、计算复杂度高，难以直接部署到边缘设备（如车载计算单元、树莓派等）上运行。因此，模型轻量化部署成为解决这一问题的关键。

轻量化部署的目标是在保证模型识别精度的前提下，显著降低模型的计算量（FLOPs）和存储需求（参数量）。这不仅能够提升模型的推理速度，还能减少内存占用，使其更适合在资源受限的边缘设备上运行。

在模型轻量化领域，常见的优化技术包括模型剪枝（Pruning）、量化（Quantization）和知识蒸馏（Knowledge Distillation）。以下是对这些技术的优缺点分析：

模型剪枝
优点：通过移除冗余的神经元或通道，显著减少模型参数量和计算量。
缺点：剪枝后的模型可能需要微调以恢复精度，且剪枝策略的设计较为复杂。
量化
优点：将模型权重和激活值从浮点数（FP32）转换为低精度整数（如 INT8），减少存储和计算开销。
缺点：量化可能导致精度下降，尤其是在低比特量化时。
知识蒸馏
优点：通过教师 - 学生模型的方式，将大模型的知识迁移到小模型中，提升小模型的性能。
缺点：训练过程复杂，需要额外的教师模型。

综合考虑，本文选择 通道剪枝 和动态量化 作为轻量化部署的核心技术，因其实现简单且效果显著。

通道剪枝是一种结构化剪枝方法，通过移除卷积层中不重要的通道来减少模型的计算量。以下是使用 PyTorch 实现通道剪枝的关键步骤：

重要性评估：计算每个通道的 L1 范数，范数较小的通道被认为是不重要的。
剪枝掩码生成：根据设定的剪枝比例（如 30%），生成剪枝掩码。
模型微调：对剪枝后的模型进行微调，以恢复精度。

以下是代码示例：

import torch
import torch.nn as nn

def channel_pruning(model, pruning_rate=0.3):
    for name, module in model.named_modules():
        if isinstance(module, nn.Conv2d):
            weights = module.weight.data
            # 计算每个通道的 L1 范数
            channel_norms = torch.norm(weights, p=1, dim=(1, 2, 3))
            # 确定剪枝阈值
            threshold = torch.quantile(channel_norms, pruning_rate)
            # 生成剪枝掩码
            mask = channel_norms > threshold
            # 应用剪枝
            pruned_weights = weights[mask, :, :, :]
            new_conv = nn.Conv2d(in_channels=int(mask.sum()),
                out_channels=module.out_channels,
                kernel_size=module.kernel_size,
                stride=module.stride,
                padding=module.padding,
                bias=module.bias is not None
            )
            new_conv.weight.data = pruned_weights
            if module.bias is not None:
                new_conv.bias.data = module.bias.data
            # 替换原始卷积层
            parent_name = name.rsplit('.', 1)[0]
            child_name = name.rsplit('.', 1)[1]
            parent_module = model.get_submodule(parent_name)
            setattr(parent_module, child_name, new_conv)
    return model

动态量化是一种在推理过程中动态计算量化参数的技术，适用于 PyTorch 模型。以下是实现动态量化的步骤：

模型准备：确保模型的所有层都支持量化。
量化配置：选择量化类型（如 INT8）和量化范围（动态或静态）。
量化模型 ：使用 PyTorch 的torch.quantization.quantize_dynamic 函数对模型进行量化。

以下是代码示例：

import torch.quantization

# 原始模型
model = ...  # 加载训练好的模型
model.eval()

# 动态量化
quantized_model = torch.quantization.quantize_dynamic(
    model,  # 原始模型
    {torch.nn.Linear, torch.nn.Conv2d},  # 需要量化的层类型
    dtype=torch.qint8  # 量化类型
)

以下是一个完整的 PyTorch 实现，包含通道剪枝和动态量化：

import torch
import torch.nn as nn
import torch.quantization

# 定义交通标志识别模型（示例）class TrafficSignModel(nn.Module):
    def __init__(self):
        super(TrafficSignModel, self).__init__()
        self.conv1 = nn.Conv2d(3, 32, kernel_size=3, stride=1, padding=1)
        self.conv2 = nn.Conv2d(32, 64, kernel_size=3, stride=1, padding=1)
        self.fc = nn.Linear(64 * 32 * 32, 10)  # 假设输出 10 类

    def forward(self, x):
        x = torch.relu(self.conv1(x))
        x = torch.max_pool2d(x, 2)
        x = torch.relu(self.conv2(x))
        x = torch.max_pool2d(x, 2)
        x = x.view(x.size(0), -1)
        x = self.fc(x)
        return x

# 通道剪枝函数
def channel_pruning(model, pruning_rate=0.3):
    for name, module in model.named_modules():
        if isinstance(module, nn.Conv2d):
            weights = module.weight.data
            channel_norms = torch.norm(weights, p=1, dim=(1, 2, 3))
            threshold = torch.quantile(channel_norms, pruning_rate)
            mask = channel_norms > threshold
            pruned_weights = weights[mask, :, :, :]
            new_conv = nn.Conv2d(in_channels=int(mask.sum()),
                out_channels=module.out_channels,
                kernel_size=module.kernel_size,
                stride=module.stride,
                padding=module.padding,
                bias=module.bias is not None
            )
            new_conv.weight.data = pruned_weights
            if module.bias is not None:
                new_conv.bias.data = module.bias.data
            parent_name = name.rsplit('.', 1)[0]
            child_name = name.rsplit('.', 1)[1]
            parent_module = model.get_submodule(parent_name)
            setattr(parent_module, child_name, new_conv)
    return model

# 加载模型
model = TrafficSignModel()
model.load_state_dict(torch.load('traffic_sign_model.pth'))
model.eval()

# 通道剪枝
pruned_model = channel_pruning(model)

# 动态量化
quantized_model = torch.quantization.quantize_dynamic(
    pruned_model,
    {torch.nn.Linear, torch.nn.Conv2d},
    dtype=torch.qint8
)

# 保存量化模型
torch.save(quantized_model.state_dict(), 'quantized_traffic_sign_model.pth')

在树莓派 4B（4GB 内存）上对原始模型、剪枝后模型和量化后模型进行性能测试，结果如下：