#!/bin/bash
# 一键启动模型分析

set -e

cd /workspace/bevfusion

echo "========================================================================"
echo "BEVFusion模型优化 - 阶段1: 模型分析"
echo "========================================================================"
echo ""
echo "Baseline: Epoch 23"
echo "  - NDS: 0.6941"
echo "  - mAP: 0.6446"
echo "  - mIoU: 0.4130"
echo ""
echo "目标: 分析模型结构，确定优化策略"
echo ""

# 配置文件和checkpoint
CONFIG="configs/nuscenes/det/transfusion/secfpn/camera+lidar/swint_v0p075/multitask_enhanced_phase1_HIGHRES.yaml"
CHECKPOINT="runs/enhanced_from_epoch19/epoch_23.pth"

# 创建输出目录
TIMESTAMP=$(date +%Y%m%d_%H%M%S)
OUTPUT_DIR="analysis_results/analysis_${TIMESTAMP}"
mkdir -p "$OUTPUT_DIR"

echo "配置文件: $CONFIG"
echo "Checkpoint: $CHECKPOINT"
echo "输出目录: $OUTPUT_DIR"
echo ""

# 检查文件是否存在
if [ ! -f "$CONFIG" ]; then
    echo "错误: 配置文件不存在: $CONFIG"
    exit 1
fi

if [ ! -f "$CHECKPOINT" ]; then
    echo "错误: Checkpoint不存在: $CHECKPOINT"
    exit 1
fi

# 1. 模型复杂度分析
echo "========== 1. 模型复杂度分析 =========="
echo "分析参数量、模型大小、各模块分布..."
echo ""

# 设置Python路径
export PATH=/opt/conda/bin:$PATH
export PYTHONPATH=/workspace/bevfusion:$PYTHONPATH

/opt/conda/bin/python tools/analysis/model_complexity_simple.py \
    "$CONFIG" \
    "$CHECKPOINT" \
    2>&1 | tee "$OUTPUT_DIR/model_complexity.txt"

echo ""
echo "✅ 模型复杂度分析完成"
echo "   结果保存到: $OUTPUT_DIR/model_complexity.txt"
echo ""

# 2. 提取关键数据
echo "========== 2. 提取关键数据 =========="

# 从分析结果中提取关键信息
TOTAL_PARAMS=$(grep "总参数量" "$OUTPUT_DIR/model_complexity.txt" | grep -oP '\d+\.\d+M' | head -1)
MODEL_SIZE_FP32=$(grep "模型大小(FP32)" "$OUTPUT_DIR/model_complexity.txt" | grep -oP '\d+\.\d+ MB')
MODEL_SIZE_INT8=$(grep "模型大小(INT8)" "$OUTPUT_DIR/model_complexity.txt" | grep -oP '\d+\.\d+ MB')

echo "总参数量: $TOTAL_PARAMS"
echo "FP32大小: $MODEL_SIZE_FP32"
echo "INT8大小: $MODEL_SIZE_INT8"
echo ""

# 3. 生成优化建议摘要
echo "========== 3. 生成优化建议 =========="

cat > "$OUTPUT_DIR/optimization_recommendations.md" << 'EOF'
# BEVFusion模型优化建议

**生成时间**: $(date)
**Baseline**: Epoch 23 (NDS 0.6941, mAP 0.6446, mIoU 0.4130)

---

## 📊 模型现状

EOF

# 添加分析结果
echo "### 参数统计" >> "$OUTPUT_DIR/optimization_recommendations.md"
echo '```' >> "$OUTPUT_DIR/optimization_recommendations.md"
grep -A 5 "总体统计" "$OUTPUT_DIR/model_complexity.txt" | tail -5 >> "$OUTPUT_DIR/optimization_recommendations.md"
echo '```' >> "$OUTPUT_DIR/optimization_recommendations.md"
echo "" >> "$OUTPUT_DIR/optimization_recommendations.md"

echo "### 各模块分布" >> "$OUTPUT_DIR/optimization_recommendations.md"
echo '```' >> "$OUTPUT_DIR/optimization_recommendations.md"
grep -A 20 "各模块参数统计" "$OUTPUT_DIR/model_complexity.txt" | tail -15 >> "$OUTPUT_DIR/optimization_recommendations.md"
echo '```' >> "$OUTPUT_DIR/optimization_recommendations.md"
echo "" >> "$OUTPUT_DIR/optimization_recommendations.md"

cat >> "$OUTPUT_DIR/optimization_recommendations.md" << 'EOF'

---

## 🎯 优化策略

### 阶段1: 结构化剪枝 (3-5天)

**目标**: 参数量减少40-45%

**策略**:
1. **Camera Encoder (SwinTransformer)** - 优先剪枝
   - 通道剪枝: 减少25-30% channels
   - 预期减少: ~15M参数
   
2. **Decoder** - 次要剪枝
   - 通道剪枝: 减少20-25% channels
   - 预期减少: ~4M参数
   
3. **FPN Neck** - 适度剪枝
   - 通道剪枝: 减少20% channels
   - 预期减少: ~8M参数

4. **Heads** - 谨慎剪枝
   - 通道剪枝: 减少10-15% channels
   - 预期减少: ~3M参数

**预期结果**:
- 参数量: 110M → 60-65M
- 模型大小: 441MB → 240-260MB
- 推理加速: 20-30%
- 精度损失: <1.5%

---

### 阶段2: INT8量化 (4-6天)

**目标**: 模型大小减少75%，推理加速2-3倍

**策略**:
1. **PTQ (Post-Training Quantization)** - 快速验证
   - 使用100-200个校准样本
   - 预期精度损失: 2-3%

2. **QAT (Quantization-Aware Training)** - 精度恢复
   - 训练5个epochs
   - 学习率: 1e-6
   - 预期精度恢复: 1-2%

**预期结果**:
- 模型大小: 240MB → 60MB (FP32→INT8)
- 推理加速: 2-3倍
- 累计精度损失: <3%

---

## 📋 实施计划

### Day 1: 模型分析 ✅ (今天)
- [x] 模型复杂度分析
- [ ] 确定剪枝目标层
- [ ] 准备剪枝工具

### Day 2-4: 剪枝实施
- [ ] 实施结构化剪枝
- [ ] 剪枝模型微调 (3 epochs)
- [ ] 评估剪枝效果

### Day 5: PTQ量化
- [ ] PTQ量化
- [ ] 精度评估
- [ ] 决定是否需要QAT

### Day 6-8: QAT训练
- [ ] QAT训练 (5 epochs)
- [ ] 最终评估
- [ ] 生成优化报告

---

## 🎯 成功标准

### 必须达到
- ✅ 参数量 < 70M
- ✅ 模型大小 < 70MB (INT8)
- ✅ NDS > 0.67 (精度损失 < 3%)
- ✅ mAP > 0.62 (精度损失 < 4%)

### 理想目标
- 🌟 参数量 < 60M
- 🌟 模型大小 < 60MB (INT8)
- 🌟 NDS > 0.68 (精度损失 < 2%)
- 🌟 推理加速 > 2.5倍

---

**下一步**: 准备剪枝工具和配置

EOF

echo "✅ 优化建议已生成"
echo "   文件: $OUTPUT_DIR/optimization_recommendations.md"
echo ""

# 4. 显示摘要
echo "========== 4. 分析摘要 =========="
echo ""
cat "$OUTPUT_DIR/optimization_recommendations.md" | head -50
echo ""
echo "... (完整内容见文件)"
echo ""

# 5. 总结
echo "========================================================================"
echo "分析完成！"
echo "========================================================================"
echo ""
echo "输出目录: $OUTPUT_DIR"
echo ""
echo "生成的文件:"
echo "  - model_complexity.txt          # 详细分析结果"
echo "  - optimization_recommendations.md  # 优化建议"
echo ""
echo "下一步:"
echo "  1. 查看分析结果: cat $OUTPUT_DIR/optimization_recommendations.md"
echo "  2. 准备剪枝工具: 安装torch-pruning"
echo "  3. 开始剪枝实施"
echo ""
echo "快速查看命令:"
echo "  cat $OUTPUT_DIR/model_complexity.txt | grep -A 10 '总体统计'"
echo "  cat $OUTPUT_DIR/model_complexity.txt | grep -A 20 '各模块参数统计'"
echo "  cat $OUTPUT_DIR/model_complexity.txt | grep -A 10 '优化建议'"
echo ""
echo "========================================================================"