#模型轻量化：MobileNet、量化、剪枝与边缘部署详解

"""
模型轻量化的重要性：

1. 设备限制：
   - 移动设备：CPU/GPU资源有限
   - IoT设备：计算能力和内存受限
   - 嵌入式系统：功耗和存储限制

2. 实时性要求：
   - 低延迟推理
   - 高吞吐量处理
   - 节省带宽成本

3. 成本效益：
   - 降低硬件成本
   - 减少云服务费用
   - 提高能源效率
"""

def lightweight_importance():
    """
    模型轻量化的重要性
    """
    importance_factors = {
        "边缘计算": "在设备端进行推理，保护隐私",
        "实时处理": "满足低延迟应用需求",
        "成本控制": "降低硬件和云服务成本",
        "能效优化": "延长电池续航时间",
        "隐私保护": "减少数据传输和云端处理"
    }
    
    print("模型轻量化的重要性:")
    for factor, desc in importance_factors.items():
        print(f"• {factor}: {desc}")

lightweight_importance()

def evaluation_metrics():
    """
    模型轻量化评估指标
    """
    metrics = {
        "参数量 (Params)": "模型参数的总数，影响模型大小",
        "计算量 (FLOPs)": "浮点运算次数，影响推理速度",
        "内存占用": "推理过程中的内存使用量",
        "推理延迟": "单次推理所需时间",
        "能耗": "推理过程的功耗消耗",
        "准确率": "模型性能的保持程度"
    }
    
    print("模型轻量化评估指标:")
    for metric, desc in metrics.items():
        print(f"• {metric}: {desc}")

evaluation_metrics()

import torch
import torch.nn as nn
import torch.nn.functional as F

class DepthwiseSeparableConv(nn.Module):
    """
    深度可分离卷积实现
    """
    def __init__(self, in_channels, out_channels, kernel_size=3, stride=1, padding=1):
        super(DepthwiseSeparableConv, self).__init__()
        
        # 深度卷积：每个输入通道单独卷积
        self.depthwise = nn.Conv2d(
            in_channels, in_channels, 
            kernel_size=kernel_size, stride=stride, 
            padding=padding, groups=in_channels, bias=False
        )
        
        # 点卷积：1x1卷积融合通道信息
        self.pointwise = nn.Conv2d(
            in_channels, out_channels, 
            kernel_size=1, stride=1, padding=0, bias=False
        )
        
        self.bn1 = nn.BatchNorm2d(in_channels)
        self.bn2 = nn.BatchNorm2d(out_channels)
        self.relu = nn.ReLU(inplace=True)
    
    def forward(self, x):
        x = self.depthwise(x)
        x = self.bn1(x)
        x = self.relu(x)
        x = self.pointwise(x)
        x = self.bn2(x)
        x = self.relu(x)
        return x

def depthwise_separable_explanation():
    """
    深度可分离卷积解释
    """
    print("深度可分离卷积:")
    print("• 普通卷积: (K×K×C_in×C_out) × H×W")
    print("• 深度卷积: (K×K×C_in) × H×W")
    print("• 点卷积: (1×1×C_in×C_out) × H×W")
    print("• 计算量减少: 约 8-9 倍")

depthwise_separable_explanation()

class MobileNetBlock(nn.Module):
    """
    MobileNet基本块
    """
    def __init__(self, in_channels, out_channels, stride=1):
        super(MobileNetBlock, self).__init__()
        
        self.conv = DepthwiseSeparableConv(in_channels, out_channels, stride=stride)
    
    def forward(self, x):
        return self.conv(x)

class MobileNetV1(nn.Module):
    """
    MobileNetV1实现
    """
    def __init__(self, num_classes=1000, width_multiplier=1.0):
        super(MobileNetV1, self).__init__()
        
        # 调整通道数
        def make_divisible(v, divisor=8, min_value=None):
            if min_value is None:
                min_value = divisor
            new_v = max(min_value, int(v + divisor / 2) // divisor * divisor)
            if new_v < 0.9 * v:
                new_v += divisor
            return new_v
        
        input_channel = make_divisible(32 * width_multiplier)
        
        # 基础卷积层
        self.first_conv = nn.Conv2d(3, input_channel, 3, 2, 1, bias=False)
        self.bn1 = nn.BatchNorm2d(input_channel)
        self.relu = nn.ReLU(inplace=True)
        
        # MobileNet配置
        layers_config = [
            # [out_channels, repeats, stride]
            [64, 1, 1],
            [128, 2, 2],
            [256, 2, 2],
            [512, 6, 2],
            [1024, 2, 2]
        ]
        
        layers = []
        for out_channels, repeats, stride in layers_config:
            out_channels = make_divisible(out_channels * width_multiplier)
            layers.append(MobileNetBlock(input_channel, out_channels, stride))
            input_channel = out_channels
            
            for _ in range(1, repeats):
                layers.append(MobileNetBlock(input_channel, out_channels, 1))
        
        self.features = nn.Sequential(*layers)
        
        # 全局平均池化和分类层
        self.avgpool = nn.AdaptiveAvgPool2d((1, 1))
        self.classifier = nn.Linear(input_channel, num_classes)
        
        # 初始化权重
        for m in self.modules():
            if isinstance(m, nn.Conv2d):
                nn.init.kaiming_normal_(m.weight, mode='fan_out', nonlinearity='relu')
            elif isinstance(m, (nn.BatchNorm2d, nn.GroupNorm)):
                nn.init.constant_(m.weight, 1)
                nn.init.constant_(m.bias, 0)
    
    def forward(self, x):
        x = self.first_conv(x)
        x = self.bn1(x)
        x = self.relu(x)
        
        x = self.features(x)
        
        x = self.avgpool(x)
        x = torch.flatten(x, 1)
        x = self.classifier(x)
        
        return x

def mobilenet_comparison():
    """
    MobileNet与其他模型参数量对比
    """
    print("模型参数量对比:")
    print("• ResNet-50: 25.6M 参数")
    print("• MobileNetV1: 4.2M 参数")
    print("• MobileNetV2: 3.5M 参数")
    print("• MobileNetV3: 5.4M 参数")
    print("• 计算量减少约 80-90%")

mobilenet_comparison()

class InvertedResidual(nn.Module):
    """
    MobileNetV2的倒残差块
    """
    def __init__(self, in_channels, out_channels, stride, expand_ratio):
        super(InvertedResidual, self).__init__()
        self.stride = stride
        hidden_dim = int(round(in_channels * expand_ratio))
        
        # 是否使用残差连接
        self.use_res_connect = self.stride == 1 and in_channels == out_channels
        
        layers = []
        if expand_ratio != 1:
            # 扩展层
            layers.append(nn.Conv2d(in_channels, hidden_dim, 1, 1, 0, bias=False))
            layers.append(nn.BatchNorm2d(hidden_dim))
            layers.append(nn.ReLU6(inplace=True))
        
        # 深度卷积
        layers.append(nn.Conv2d(hidden_dim, hidden_dim, 3, stride, 1, groups=hidden_dim, bias=False))
        layers.append(nn.BatchNorm2d(hidden_dim))
        layers.append(nn.ReLU6(inplace=True))
        
        # 压缩层
        layers.append(nn.Conv2d(hidden_dim, out_channels, 1, 1, 0, bias=False))
        layers.append(nn.BatchNorm2d(out_channels))
        
        self.conv = nn.Sequential(*layers)
    
    def forward(self, x):
        if self.use_res_connect:
            return x + self.conv(x)
        else:
            return self.conv(x)

class MobileNetV2(nn.Module):
    """
    MobileNetV2实现
    """
    def __init__(self, num_classes=1000, width_mult=1.0):
        super(MobileNetV2, self).__init__()
        input_channel = int(32 * width_mult)
        last_channel = int(1280 * max(1.0, width_mult))
        
        # 基础卷积层
        self.features = [nn.Conv2d(3, input_channel, 3, 2, 1, bias=False),
                         nn.BatchNorm2d(input_channel),
                         nn.ReLU6(inplace=True)]
        
        # MobileNetV2配置
        inverted_residual_setting = [
            # t, c, n, s
            [1, 16, 1, 1],    # stride 2 -> 1
            [6, 24, 2, 2],
            [6, 32, 3, 2],
            [6, 64, 4, 2],
            [6, 96, 3, 1],
            [6, 160, 3, 2],
            [6, 320, 1, 1],
        ]
        
        for t, c, n, s in inverted_residual_setting:
            output_channel = int(c * width_mult)
            for i in range(n):
                stride = s if i == 0 else 1
                self.features.append(InvertedResidual(input_channel, output_channel, stride, expand_ratio=t))
                input_channel = output_channel
        
        # 最后一层
        self.features.append(nn.Conv2d(input_channel, last_channel, 1, 1, 0, bias=False))
        self.features.append(nn.BatchNorm2d(last_channel))
        self.features.append(nn.ReLU6(inplace=True))
        
        self.features = nn.Sequential(*self.features)
        
        # 分类器
        self.avgpool = nn.AdaptiveAvgPool2d((1, 1))
        self.classifier = nn.Sequential(
            nn.Dropout(0.2),
            nn.Linear(last_channel, num_classes),
        )
    
    def forward(self, x):
        x = self.features(x)
        x = self.avgpool(x)
        x = torch.flatten(x, 1)
        x = self.classifier(x)
        return x

def mobilenet_v2_improvements():
    """
    MobileNetV2改进点
    """
    improvements = [
        "倒残差结构: 先扩展再压缩",
        "线性瓶颈: 最后一层不用激活函数",
        "残差连接: 提高梯度流动",
        "ReLU6激活: 适合移动端量化"
    ]
    
    print("MobileNetV2改进:")
    for improvement in improvements:
        print(f"• {improvement}")

mobilenet_v2_improvements()

def quantization_basics():
    """
    模型量化基础概念
    """
    print("模型量化基础:")
    print("• 32位浮点数 -> 8位整数")
    print("• 模型大小减少 4 倍")
    print("• 推理速度提升 2-3 倍")
    print("• 精度损失可控")

def quantization_types():
    """
    量化类型
    """
    types = {
        "静态量化": "训练后量化，需要校准数据",
        "动态量化": "运行时量化，权重静态量化",
        "量化感知训练": "训练过程中模拟量化效果"
    }
    
    print("量化类型:")
    for qtype, desc in types.items():
        print(f"• {qtype}: {desc}")

quantization_types()

import torch.quantization as quantization
from torch.quantization import QuantStub, DeQuantStub

class QuantizableMobileNetV2(nn.Module):
    """
    可量化的MobileNetV2
    """
    def __init__(self, num_classes=1000):
        super(QuantizableMobileNetV2, self).__init__()
        
        # 量化桩
        self.quant = QuantStub()
        self.dequant = DeQuantStub()
        
        # 原始MobileNetV2结构
        self.model = MobileNetV2(num_classes=num_classes)
    
    def forward(self, x):
        x = self.quant(x)
        x = self.model(x)
        x = self.dequant(x)
        return x
    
    def fuse_model(self):
        """
        融合BN层到卷积层
        """
        torch.quantization.fuse_modules(self.model, [['first_conv', 'bn1']], inplace=True)

def static_quantization_example():
    """
    静态量化示例
    """
    print("静态量化流程:")
    print("""
# 1. 准备模型
model = QuantizableMobileNetV2()
model.eval()

# 2. 融合BN层
model.fuse_model()

# 3. 设置量化配置
model.qconfig = torch.quantization.get_default_qconfig('fbgemm')

# 4. 准备量化
torch.quantization.prepare(model, inplace=True)

# 5. 校准（使用少量数据）
for data, target in calibration_loader:
    model(data)

# 6. 转换为量化模型
torch.quantization.convert(model, inplace=True)
""")

static_quantization_example()

def quantization_aware_training():
    """
    量化感知训练示例
    """
    print("量化感知训练:")
    print("""
# 1. 设置QAT配置
model = QuantizableMobileNetV2()
model.train()
model.qconfig = torch.quantization.get_default_qat_qconfig('fbgemm')

# 2. 准备QAT
model = torch.quantization.prepare_qat(model, inplace=True)

# 3. 正常训练
optimizer = torch.optim.Adam(model.parameters(), lr=0.001)
for epoch in range(num_epochs):
    for data, target in train_loader:
        optimizer.zero_grad()
        output = model(data)
        loss = criterion(output, target)
        loss.backward()
        optimizer.step()

# 4. 转换为推理模型
model.eval()
torch.quantization.convert(model, inplace=True)
""")

quantization_aware_training()

import torch.nn.utils.prune as prune

def pruning_concepts():
    """
    模型剪枝概念
    """
    concepts = {
        "非结构化剪枝": "移除单个权重，保持稠密矩阵结构",
        "结构化剪枝": "移除整个通道、滤波器或层",
        "幅度剪枝": "基于权重绝对值大小剪枝",
        "迭代剪枝": "逐步增加剪枝率"
    }
    
    print("模型剪枝概念:")
    for concept, desc in concepts.items():
        print(f"• {concept}: {desc}")

pruning_concepts()

class PruningExample(nn.Module):
    """
    剪枝示例模型
    """
    def __init__(self):
        super(PruningExample, self).__init__()
        self.conv1 = nn.Conv2d(3, 64, 3, padding=1)
        self.conv2 = nn.Conv2d(64, 128, 3, padding=1)
        self.fc1 = nn.Linear(128, 10)
    
    def forward(self, x):
        x = F.relu(self.conv1(x))
        x = F.relu(self.conv2(x))
        x = F.adaptive_avg_pool2d(x, (1, 1))
        x = x.view(x.size(0), -1)
        x = self.fc1(x)
        return x

def pruning_examples():
    """
    剪枝示例
    """
    print("模型剪枝示例:")
    print("""
import torch.nn.utils.prune as prune

model = PruningExample()

# 1. 非结构化剪枝 - L1范数
prune.l1_unstructured(model.conv1, name='weight', amount=0.2)  # 剪掉20%权重

# 2. 结构化剪枝 - 移除整个输出通道
prune.ln_structured(model.conv1, name='weight', amount=32, n=2, dim=0)  # 移除32个输出通道

# 3. 自定义剪枝 - 基于自定义掩码
custom_mask = torch.randn_like(model.conv1.weight)
zero_mask = custom_mask.abs() < 0.1
prune.custom_from_mask(model.conv1, name='weight', mask=zero_mask)

# 4. 移除剪枝重新参数化
prune.remove(model.conv1, 'weight')  # 移除剪枝，保留剪枝后的权重
""")

pruning_examples()

def iterative_pruning_strategy():
    """
    迭代剪枝策略
    """
    print("迭代剪枝流程:")
    print("""
def iterative_pruning(model, final_sparsity=0.5, num_iterations=10):
    # 计算每次迭代的剪枝率
    sparsity_step = final_sparsity / num_iterations
    
    for iteration in range(num_iterations):
        # 计算当前目标稀疏度
        current_sparsity = (iteration + 1) * sparsity_step
        
        # 对每一层进行剪枝
        for name, module in model.named_modules():
            if isinstance(module, nn.Conv2d) or isinstance(module, nn.Linear):
                prune.l1_unstructured(module, name='weight', amount=current_sparsity)
        
        # 微调模型
        fine_tune(model)
    
    return model
""")

iterative_pruning_strategy()

class TeacherStudent(nn.Module):
    """
    教师-学生网络示例
    """
    def __init__(self, teacher, student):
        super(TeacherStudent, self).__init__()
        self.teacher = teacher
        self.student = student
    
    def forward(self, x):
        with torch.no_grad():
            teacher_output = self.teacher(x)
        student_output = self.student(x)
        return teacher_output, student_output

def knowledge_distillation_basics():
    """
    知识蒸馏基础
    """
    print("知识蒸馏:")
    print("• 教师模型: 大而准确的模型")
    print("• 学生模型: 小而快速的模型")
    print("• 软标签: 教师模型的概率分布")
    print("• 温度参数: 控制概率分布的平滑度")

def distillation_loss_function(teacher_output, student_output, target, temperature=4.0, alpha=0.7):
    """
    知识蒸馏损失函数
    """
    # 软目标损失
    soft_loss = F.kl_div(
        F.log_softmax(student_output / temperature, dim=1),
        F.softmax(teacher_output / temperature, dim=1),
        reduction='batchmean'
    ) * (temperature ** 2)
    
    # 硬目标损失
    hard_loss = F.cross_entropy(student_output, target)
    
    # 总损失
    total_loss = alpha * soft_loss + (1 - alpha) * hard_loss
    
    return total_loss

def knowledge_distillation_example():
    """
    知识蒸馏示例
    """
    print("知识蒸馏实现:")
    print("""
# 训练循环
for epoch in range(num_epochs):
    for data, target in train_loader:
        teacher_output, student_output = model(data)
        
        # 计算蒸馏损失
        loss = distillation_loss_function(
            teacher_output, student_output, target,
            temperature=4.0, alpha=0.7
        )
        
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()
""")

knowledge_distillation_example()

def neural_architecture_search():
    """
    神经架构搜索
    """
    print("神经架构搜索(NAS):")
    print("• 自动设计最优网络结构")
    "• MnasNet: 移动端NAS模型"
    print("• ProxylessNAS: 直接在目标硬件上搜索")
    print("• FBNet: 基于强化学习的NAS")

neural_architecture_search()

def deep_compression():
    """
    深度压缩技术
    """
    print("深度压缩:")
    print("• 剪枝: 移除不重要的连接")
    print("• 量化: 降低权重精度")
    print("• Huffman编码: 压缩权重存储")
    print("• 组合效果: 可达500倍压缩比")

deep_compression()

def deployment_frameworks():
    """
    模型部署框架对比
    """
    frameworks = {
        "TensorRT": "NVIDIA GPU优化，高性能推理",
        "OpenVINO": "Intel硬件优化，跨平台支持",
        "TensorFlow Lite": "移动端推理，支持量化",
        "ONNX Runtime": "跨框架推理，多后端支持",
        "PyTorch Mobile": "PyTorch原生移动端支持"
    }
    
    print("模型部署框架:")
    for framework, desc in frameworks.items():
        print(f"• {framework}: {desc}")

deployment_frameworks()

def optimization_strategies():
    """
    性能优化策略
    """
    strategies = [
        "选择合适的轻量化架构",
        "使用混合精度训练和推理",
        "优化数据加载和预处理",
        "使用模型并行和数据并行",
        "针对硬件特性优化实现",
        "减少内存拷贝和数据传输"
    ]
    
    print("性能优化策略:")
    for strategy in strategies:
        print(f"• {strategy}")

optimization_strategies()

def mobile_deployment_example():
    """
    移动端部署示例
    """
    print("移动端部署流程:")
    print("""
# 1. 模型训练和优化
model = MobileNetV2()
# ... 训练代码 ...

# 2. 模型量化
model.eval()
model.qconfig = torch.quantization.get_default_qconfig('qnnpack')
torch.quantization.prepare(model, inplace=True)
# 校准...
torch.quantization.convert(model, inplace=True)

# 3. 转换为TFLite格式
import torch
torchscript_model = torch.jit.trace(model, example_input)
torch.jit.save(torchscript_model, "mobile_model.pt")

# 4. 在移动应用中使用
# Android/iOS应用集成模型文件
# 实时推理处理摄像头输入
""")

mobile_deployment_example()

def edge_deployment_considerations():
    """
    边缘设备部署考虑因素
    """
    considerations = {
        "硬件限制": "CPU/GPU/内存/NPU能力",
        "功耗要求": "电池续航和散热管理", 
        "实时性": "延迟和吞吐量要求",
        "可靠性": "长时间稳定运行",
        "安全性": "模型和数据保护"
    }
    
    print("边缘部署考虑因素:")
    for factor, desc in considerations.items():
        print(f"• {factor}: {desc}")

edge_deployment_considerations()