RoRD-Layout-Recognation/docs/feature_work.md

后续工作

本文档整合了 RoRD 项目的优化待办清单和训练需求，用于规划未来的开发和实验工作。

---

RoRD 项目优化待办清单

本文档旨在为 RoRD (Rotation-Robust Descriptors) 项目提供一系列可行的优化任务。各项任务按优先级和模块划分，您可以根据项目进度和资源情况选择执行。

一、 数据策略与增强 (Data Strategy & Augmentation)

目标：提升模型的鲁棒性和泛化能力，减少对大量真实数据的依赖。

• 引入弹性变形 (Elastic Transformations)
  • ✔️ 价值: 模拟芯片制造中可能出现的微小物理形变，使模型对非刚性变化更鲁棒。
  • 📝 执行方案:
    1. 添加 albumentations 库作为项目依赖。
    2. 在 train.py 的 ICLayoutTrainingDataset 类中，集成 A.ElasticTransform 到数据增强管道中。
• 创建合成版图数据生成器
  • ✔️ 价值: 解决真实版图数据获取难、数量少的问题，通过程序化生成大量多样化的训练样本。
  • 📝 执行方案:
    1. 创建一个新脚本，例如 tools/generate_synthetic_layouts.py。
    2. 利用 gdstk 库 编写函数，程序化地生成包含不同尺寸、密度和类型标准单元的 GDSII 文件。
    3. 结合 tools/layout2png.py 的逻辑，将生成的版图批量转换为 PNG 图像，用于扩充训练集。

二、 模型架构 (Model Architecture)

目标：提升模型的特征提取效率和精度，降低计算资源消耗。

• 实验更现代的骨干网络 (Backbone)
  • ✔️ 价值: VGG-16 经典但效率偏低。新架构（如 ResNet, EfficientNet）能以更少的参数量和计算量达到更好的性能。
  • 📝 执行方案:
    1. 在 models/rord.py 中，修改 RoRD 类的 __init__ 方法。
    2. 使用 torchvision.models 替换 vgg16。可尝试 models.resnet34(pretrained=True) 或 models.efficientnet_b0(pretrained=True) 作为替代方案。
    3. 相应地调整检测头和描述子头的输入通道数。
• 集成注意力机制 (Attention Mechanism)
  • ✔️ 价值: 引导模型自动关注版图中的关键几何结构（如边角、交点），忽略大面积的空白或重复区域，提升特征质量。
  • 📝 执行方案:
    1. 寻找一个可靠的注意力模块实现，如 CBAM 或 SE-Net。
    2. 在 models/rord.py 中，将该模块插入到 self.backbone 和两个 head 之间。

三、 训练与损失函数 (Training & Loss Function)

目标：优化训练过程的稳定性，提升模型收敛效果。

• 实现损失函数的自动加权
  • ✔️ 价值: 当前检测损失和描述子损失是等权重相加，手动调参困难。自动加权可以使模型自主地平衡不同任务的优化难度。
  • 📝 执行方案:
    1. 参考学术界关于“多任务学习中的不确定性加权” (Uncertainty Weighting) 的论文。
    2. 在 train.py 中，将损失权重定义为两个可学习的参数 log_var_a 和 log_var_b。
    3. 将总损失函数修改为 loss = torch.exp(-log_var_a) * det_loss + log_var_a + torch.exp(-log_var_b) * desc_loss + log_var_b。
    4. 将这两个新参数加入到优化器中进行训练。
• 实现基于关键点响应的困难样本采样
  • ✔️ 价值: 提升描述子学习的效率。只在模型认为是“关键点”的区域进行采样，能让模型更专注于学习有区分度的特征。
  • 📝 执行方案:
    1. 在 train.py 的 compute_description_loss 函数中。
    2. 获取 det_original 的输出图，进行阈值处理或 Top-K 选择，得到关键点的位置坐标。
    3. 使用这些坐标，而不是 torch.linspace 生成的网格坐标，作为采样点来提取 anchor、positive 和 negative 描述子。

四、 推理与匹配 (Inference & Matching)

目标：大幅提升大尺寸版图的匹配速度和多尺度检测能力。

• 将模型改造为特征金字塔网络 (FPN) 架构
  • ✔️ 价值: 当前的多尺度匹配需要多次缩放图像并推理，速度慢。FPN 只需一次推理即可获得所有尺度的特征，极大加速匹配过程。
  • 📝 执行方案:
    1. 修改 models/rord.py，从骨干网络的不同层级（如 VGG 的 relu2_2, relu3_3, relu4_3）提取特征图。
    2. 添加上采样和横向连接层来融合这些特征图，构建出特征金字塔。
    3. 修改 match.py，使其能够直接从 FPN 的不同层级获取特征，替代原有的图像金字塔循环。
• 在滑动窗口匹配后增加关键点去重
  • ✔️ 价值: match.py 中的滑动窗口在重叠区域会产生大量重复的关键点，增加后续匹配的计算量并可能影响精度。
  • 📝 执行方案:
    1. 在 match.py 的 extract_features_sliding_window 函数返回前。
    2. 实现一个非极大值抑制 (NMS) 算法。
    3. 根据关键点的位置和检测分数（需要模型输出强度图），对 all_kps 和 all_descs 进行过滤，去除冗余点。

五、 代码与项目结构 (Code & Project Structure)

目标：提升项目的可维护性、可扩展性和易用性。

• 迁移配置到 YAML 文件
  • ✔️ 价值: config.py 不利于管理多组实验配置。YAML 文件能让每组实验的参数独立、清晰，便于复现。
  • 📝 执行方案:
    1. 创建一个 configs 目录，并编写一个 base_config.yaml 文件。
    2. 引入 OmegaConf 或 Hydra 库。
    3. 修改 train.py 和 match.py 等脚本，使其从 YAML 文件加载配置，而不是从 config.py 导入。
• 代码模块解耦
  • ✔️ 价值: train.py 文件过长，职责过多。解耦能使代码结构更清晰，符合单一职责原则。
  • 📝 执行方案:
    1. 将 ICLayoutTrainingDataset 类从 train.py 移动到 data/ic_dataset.py。
    2. 创建一个新文件 losses.py，将 compute_detection_loss 和 compute_description_loss 函数移入其中。

六、 实验跟踪与评估 (Experiment Tracking & Evaluation)

目标：建立科学的实验流程，提供更全面的模型性能度量。

• 集成实验跟踪工具 (TensorBoard / W&B)
  • ✔️ 价值: 日志文件不利于直观对比实验结果。可视化工具可以实时监控、比较多组实验的损失和评估指标。
  • 📝 执行方案:
    1. 在 train.py 中，导入 torch.utils.tensorboard.SummaryWriter。
    2. 在训练循环中，使用 writer.add_scalar() 记录各项损失值。
    3. 在验证结束后，记录评估指标和学习率等信息。
• 增加更全面的评估指标
  • ✔️ 价值: 当前的评估指标 主要关注检测框的重合度。增加 mAP 和几何误差评估能更全面地衡量模型性能。
  • 📝 执行方案:
    1. 在 evaluate.py 中，实现 mAP (mean Average Precision) 的计算逻辑。
    2. 在计算 IoU 匹配成功后，从 match_template_multiscale 返回的单应性矩阵 H 中，分解出旋转/平移等几何参数，并与真实变换进行比较，计算误差。

---

训练需求

1. 数据集类型

• 格式: 训练数据为PNG格式的集成电路 (IC) 版图图像。这些图像可以是二值化的黑白图，也可以是灰度图。
• 来源: 可以从 GDSII (.gds) 或 OASIS (.oas) 版图文件通过光栅化生成。
• 内容: 数据集应包含多种不同区域、不同风格的版图，以确保模型的泛化能力。
• 标注: 训练阶段无需任何人工标注。模型采用自监督学习，通过对原图进行旋转、镜像等几何变换来自动生成训练对。

2. 数据集大小

• 启动阶段 (功能验证): 100 - 200 张 高分辨率 (例如：2048x2048) 的版图图像。这个规模足以验证训练流程是否能跑通，损失函数是否收敛。
• 初步可用模型: 1,000 - 2,000 张 版图图像。在这个数量级上，模型能学习到比较鲁棒的几何特征，在与训练数据相似的版图上取得不错的效果。
• 生产级模型: 5,000 - 10,000+ 张 版图图像。要让模型在各种不同工艺、设计风格的版图上都具有良好的泛化能力，需要大规模、多样化的数据集。

训练脚本 train.py 会将提供的数据集自动按 80/20 的比例划分为训练集和验证集。

3. 计算资源

• 硬件: 一块支持 CUDA 的 NVIDIA GPU 是必需的。考虑到模型的 VGG-16 骨干网络和复杂的几何感知损失函数，使用中高端 GPU 会显著提升训练效率。
• 推荐型号:
  • 入门级: NVIDIA RTX 3060 / 4060
  • 主流级: NVIDIA RTX 3080 / 4070 / A4000
  • 专业级: NVIDIA RTX 3090 / 4090 / A6000
• CPU 与内存: 建议至少 8 核 CPU 和 32 GB 内存，以确保数据预处理和加载不会成为瓶颈。

4. 显存大小 (VRAM)

根据配置文件 config.py 和 train.py 中的参数，可以估算所需显存：

• 模型架构: 基于 VGG-16。
• 批次大小 (Batch Size): 默认为 8。
• 图像块大小 (Patch Size): 256x256。

综合以上参数，并考虑到梯度和优化器状态的存储开销，建议至少需要 12 GB 显存。如果显存不足，需要将 BATCH_SIZE 减小 (例如 4 或 2)，但这会牺牲训练速度和稳定性。

5. 训练时间估算

假设使用一块 NVIDIA RTX 3080 (10GB) 显卡和 2,000 张 版图图像的数据集：

• 单个 Epoch 时间: 约 15 - 25 分钟。
• 总训练时间: 配置文件中设置的总轮数 (Epochs) 为 50。
  • 50 epochs * 20 分钟/epoch ≈ 16.7 小时
• 收敛时间: 项目引入了早停机制 (patience=10)，如果验证集损失在 10 个 epoch 内没有改善，训练会提前停止。因此，实际训练时间可能在 10 到 20 小时 之间。

6. 逐步调优时间

调优是一个迭代过程，非常耗时。根据 TRAINING_STRATEGY_ANALYSIS.md 文件中提到的优化点 和进一步优化建议，调优阶段可能包括：

• 数据增强策略探索 (1-2周): 调整尺度抖动范围、亮度和对比度参数，尝试不同的噪声类型等。
• 损失函数权重平衡 (1-2周): loss_function.md 中提到了多种损失分量（BCE, SmoothL1, Triplet, Manhattan, Sparsity, Binary），调整它们之间的权重对模型性能至关重要。
• 超参数搜索 (2-4周): 对学习率、批次大小、优化器类型 (Adam, SGD等)、学习率调度策略等进行网格搜索或贝叶斯优化。
• 模型架构微调 (可选，2-4周): 尝试不同的骨干网络 (如 ResNet)、修改检测头和描述子头的层数或通道数。

总计，要达到一个稳定、可靠、泛化能力强的生产级模型，从数据准备到最终调优完成，预计需要 1 个半到 3 个月的时间。