Lung Nodule Detection with Deep Learning & RK3588S Deployment

🎯 目标：构建一个端到端的肺结节检测系统，并通过模型优化与边缘设备（RK3588S）部署，完成性能评估与对比。

---

📌 背景

• 本项目的完整的端到端方案规划及模型架构设计参考 《Deep Learning With PyTorch_Code》—— Eli Stevens, Luca Antiga, Thomas Viehmann（2021，Manning Publications） 一书。
• 本项目所用到的数据全部来自于 LUNA2016 Dataset
• 扩展部分主要体现在对模型优化后，部署到RK3588S开发板，并输出性能对比报告。

🧠 项目亮点

• 基于 PyTorch 实现 CT 影像肺结节检测
• 使用量化 (PTQ / QTA)、剪枝优化技术，提升推理速度
• 成功部署至 RK3588S，进行真实设备测试
• 全流程自动化，从数据处理到部署测试
• 附带对比图表，评估准确率/速度/内存消耗

---

📁 项目结构

├── README.md # 项目说明文件
├── data # 数据集和保存的模型
│   ├── luna # 数据集
│   └── models # 已训练，已优化的模型
├── deployment # 部署相关
│   ├── convert_rknn.py
│   ├── export_onnx.py
│   └── onnx_check.py
├── images # 图片资源
├── notebooks # Jupyter Notebooks 目录，包含各个步骤的实现
│   ├── 1_data_preprocessing.ipynb # 数据预处理
│   ├── 2_model_training.ipynb # 模型训练
│   ├── 3_model_deploy.ipynb # 部署到 RK3588S 开发板
│   └── 4_model_optimization.ipynb # 模型优化
├── optimization # 模型优化相关
│   ├── fx_quantization.py
│   ├── helper.py
│   └── pruning.py
├── scripts # 快捷脚本
│   ├── helper.py
│   ├── path.py
│   ├── run_cache_dataset.py
│   ├── run_nodule_analysis.py
│   ├── run_optimization_deployment.py
│   └── run_training.py
└── src # 核心源代码
    ├── app # 应用封装
    ├── core # 模型架构
    └── util # 工具函数

---

📊 性能评估对比表

完整报告见 reports/performance.md

模型版本	Weight Memory (KiB)	Reduction Ratio	Size (KiB)	Latency (us)	Reduction Ratio	Recall (%)	Recall Drop
原始模型	236.81	1.00x	1005.18	49079	1.00x	76.7%	-
剪枝后	64.75	3.66x	514.05	30171	1.63x	75.3%	↓ 1.4%
PTQ后	125.44	1.89x	561.83	31536	1.56x	79.6%	↑ 2.9%
剪枝+PTQ	40.81	5.80x	453.83	29710	1.65x	69.5%	↓ 7.2%
剪枝+QAT	40.56	5.84x	451.78	30191	1.63x	77.1%	↑ 0.4%

---

🚀 快速开始

# 安装依赖
conda env create -f environment.yaml

# rknn-toolkit2 的环境安装请参考官方文档，见附录-官方仓库链接

# 训练前缓存数据集到磁盘（可选）
python scripts/run_cache_dataset.py

# 训练分割模型和分类模型
python scripts/run_training.py

# 训练过程中的的模型会保存到data-unversioned/models/目录下

# 优化并部署模型到 RK3588S
python scripts/run_optimization_deployment.py

# 优化后的模型会导出到build/models/目录下

# 运行肺结节检测应用

# PC上推理模型
python scripts/run_nodule_analysis.py --platform pytorch --run-validation

# 连板推理模型
python scripts/run_nodule_analysis.py --platform rknn --segmentation-path data/models/seg/seg_model.rknn --target rk3588  --run-validation

# 更多参数选项
python scripts/run_xxx.py --help

📚 参考资料

• 《Deep Learning With PyTorch_Code》—— Eli Stevens, Luca Antiga, Thomas Viehmann（2021，Manning Publications）
• MIT 6.5940 • Fall 2024 • TinyML and Efficient Deep Learning Computing
• PyTorch Documentation
• Rockchip RKNPU Official Docs
• Get started with TensorBoard