Merge pull request '添加数据增强方案以及扩散生成模型的想法' (#5) from lingke-dataAugentation-20251020 into main
Reviewed-on: #5
This commit is contained in:
121
README.md
121
README.md
@@ -4,6 +4,29 @@
|
|||||||
|
|
||||||
|
|
||||||
|
|
||||||
|
## ⚡ Quick Start（含合成数据与H校验）
|
||||||
|
|
||||||
|
```bash
|
||||||
|
# 一键生成→渲染→预览→H校验→写回配置（开启合成混采与 Elastic）
|
||||||
|
uv run python tools/synth_pipeline.py \
|
||||||
|
--out_root data/synthetic \
|
||||||
|
--num 50 \
|
||||||
|
--dpi 600 \
|
||||||
|
--config configs/base_config.yaml \
|
||||||
|
--ratio 0.3 \
|
||||||
|
--enable_elastic \
|
||||||
|
--validate_h --validate_n 6
|
||||||
|
```
|
||||||
|
|
||||||
|
提示：zsh 下使用反斜杠续行时，确保每行末尾只有一个 `\` 且下一行不要粘连参数（避免如 `6uv` 这样的粘连）。
|
||||||
|
|
||||||
|
可选：为 KLayout 渲染指定图层配色/线宽/背景（示例：金属层绿色、过孔红色、黑底）
|
||||||
|
```bash
|
||||||
|
uv run python tools/layout2png.py \
|
||||||
|
--in data/synthetic/gds --out data/synthetic/png --dpi 800 \
|
||||||
|
--layermap '1/0:#00FF00,2/0:#FF0000' --line_width 2 --bgcolor '#000000'
|
||||||
|
```
|
||||||
|
|
||||||
## 📖 描述
|
## 📖 描述
|
||||||
|
|
||||||
本项目实现了 **RoRD (Rotation-Robust Descriptors)** 模型，这是一种先进的局部特征匹配方法，专用于集成电路（IC）版图的识别。
|
本项目实现了 **RoRD (Rotation-Robust Descriptors)** 模型，这是一种先进的局部特征匹配方法，专用于集成电路（IC）版图的识别。
|
||||||
@@ -70,7 +93,9 @@ RoRD-Layout-Recognation/
|
|||||||
├── match.py # 模板匹配脚本（FPN / 滑窗 + NMS）
|
├── match.py # 模板匹配脚本（FPN / 滑窗 + NMS）
|
||||||
├── tests/
|
├── tests/
|
||||||
│ ├── benchmark_fpn.py # FPN vs 滑窗性能对标
|
│ ├── benchmark_fpn.py # FPN vs 滑窗性能对标
|
||||||
│ └── benchmark_backbones.py # 多骨干 A/B 前向基准
|
│ ├── benchmark_backbones.py # 多骨干 A/B 前向基准
|
||||||
|
│ ├── benchmark_attention.py # 注意力 none/se/cbam A/B 基准
|
||||||
|
│ └── benchmark_grid.py # 三维基准：Backbone × Attention × Single/FPN
|
||||||
├── config.py # 兼容旧流程的 YAML 读取 shim
|
├── config.py # 兼容旧流程的 YAML 读取 shim
|
||||||
├── pyproject.toml
|
├── pyproject.toml
|
||||||
└── README.md
|
└── README.md
|
||||||
@@ -81,6 +106,11 @@ RoRD-Layout-Recognation/
|
|||||||
- **YAML 配置中心**：所有路径与超参数集中存放在 `configs/*.yaml`，通过 `utils.config_loader.load_config` 统一解析；CLI 的 `--config` 参数可切换实验配置，`to_absolute_path` 则保证相对路径相对配置文件解析。
|
- **YAML 配置中心**：所有路径与超参数集中存放在 `configs/*.yaml`，通过 `utils.config_loader.load_config` 统一解析；CLI 的 `--config` 参数可切换实验配置，`to_absolute_path` 则保证相对路径相对配置文件解析。
|
||||||
- **旧配置兼容**：`config.py` 现在仅作为兼容层，将 YAML 配置转换成原有的 Python 常量，便于逐步迁移历史代码。
|
- **旧配置兼容**：`config.py` 现在仅作为兼容层，将 YAML 配置转换成原有的 Python 常量，便于逐步迁移历史代码。
|
||||||
- **损失与数据解耦**：`losses.py` 汇总几何感知损失，`data/ic_dataset.py` 与 `utils/data_utils.py` 分离数据准备逻辑，便于扩展新的采样策略或损失项。
|
- **损失与数据解耦**：`losses.py` 汇总几何感知损失，`data/ic_dataset.py` 与 `utils/data_utils.py` 分离数据准备逻辑，便于扩展新的采样策略或损失项。
|
||||||
|
|
||||||
|
# 5. 运行 A/B 基准（骨干、注意力、三维网格）
|
||||||
|
PYTHONPATH=. uv run python tests/benchmark_backbones.py --device cpu --image-size 512 --runs 5
|
||||||
|
PYTHONPATH=. uv run python tests/benchmark_attention.py --device cpu --image-size 512 --runs 10 --backbone resnet34 --places backbone_high desc_head
|
||||||
|
PYTHONPATH=. uv run python tests/benchmark_grid.py --device cpu --image-size 512 --runs 3 --backbones vgg16 resnet34 efficientnet_b0 --attentions none se cbam --places backbone_high desc_head
|
||||||
- **日志体系**：`logging` 配置节配合 TensorBoard 集成，`train.py`、`evaluate.py`、`match.py` 可统一写入 `log_dir/子任务/experiment_name`。
|
- **日志体系**：`logging` 配置节配合 TensorBoard 集成，`train.py`、`evaluate.py`、`match.py` 可统一写入 `log_dir/子任务/experiment_name`。
|
||||||
- **模型配置扩展**：
|
- **模型配置扩展**：
|
||||||
- `model.backbone.name`: `vgg16 | resnet34 | efficientnet_b0`
|
- `model.backbone.name`: `vgg16 | resnet34 | efficientnet_b0`
|
||||||
@@ -350,6 +380,7 @@ uv run python match.py --config configs/base_config.yaml --no_nms \
|
|||||||
可参考以下文档与脚本复现并查看最新结果：
|
可参考以下文档与脚本复现并查看最新结果：
|
||||||
|
|
||||||
- CPU 多骨干 A/B 基准（512×512，5 次）：见 `docs/description/Performance_Benchmark.md`
|
- CPU 多骨干 A/B 基准（512×512，5 次）：见 `docs/description/Performance_Benchmark.md`
|
||||||
|
- 三维基准（Backbone × Attention × Single/FPN）：见 `docs/description/Performance_Benchmark.md` 与 `tests/benchmark_grid.py`
|
||||||
- FPN vs 滑窗对标脚本：`tests/benchmark_fpn.py`
|
- FPN vs 滑窗对标脚本：`tests/benchmark_fpn.py`
|
||||||
- 多骨干 A/B 基准脚本：`tests/benchmark_backbones.py`
|
- 多骨干 A/B 基准脚本：`tests/benchmark_backbones.py`
|
||||||
|
|
||||||
@@ -358,3 +389,91 @@ uv run python match.py --config configs/base_config.yaml --no_nms \
|
|||||||
## 📄 许可协议
|
## 📄 许可协议
|
||||||
|
|
||||||
本项目根据 [Apache License 2.0](LICENSE.txt) 授权。
|
本项目根据 [Apache License 2.0](LICENSE.txt) 授权。
|
||||||
|
|
||||||
|
---
|
||||||
|
|
||||||
|
## 🧪 合成数据一键流程与常见问题
|
||||||
|
|
||||||
|
### 一键命令
|
||||||
|
```bash
|
||||||
|
uv run python tools/generate_synthetic_layouts.py --out_dir data/synthetic/gds --num 200 --seed 42
|
||||||
|
uv run python tools/layout2png.py --in data/synthetic/gds --out data/synthetic/png --dpi 600
|
||||||
|
uv run python tools/preview_dataset.py --dir data/synthetic/png --out preview.png --n 8 --elastic
|
||||||
|
uv run python train.py --config configs/base_config.yaml
|
||||||
|
```
|
||||||
|
|
||||||
|
或使用单脚本一键执行（含配置写回）：
|
||||||
|
```bash
|
||||||
|
uv run python tools/synth_pipeline.py --out_root data/synthetic --num 200 --dpi 600 \
|
||||||
|
--config configs/base_config.yaml --ratio 0.3 --enable_elastic
|
||||||
|
```
|
||||||
|
|
||||||
|
### YAML 关键片段
|
||||||
|
```yaml
|
||||||
|
synthetic:
|
||||||
|
enabled: true
|
||||||
|
png_dir: data/synthetic/png
|
||||||
|
ratio: 0.3
|
||||||
|
|
||||||
|
augment:
|
||||||
|
elastic:
|
||||||
|
enabled: true
|
||||||
|
alpha: 40
|
||||||
|
sigma: 6
|
||||||
|
alpha_affine: 6
|
||||||
|
prob: 0.3
|
||||||
|
```
|
||||||
|
|
||||||
|
### 参数建议
|
||||||
|
- DPI：600–900；图形极细时可到 1200（注意磁盘占用与 IO）。
|
||||||
|
- ratio：数据少取 0.3–0.5；中等 0.2–0.3；数据多 0.1–0.2。
|
||||||
|
- Elastic：alpha=40, sigma=6, prob=0.3 为安全起点。
|
||||||
|
|
||||||
|
### FAQ
|
||||||
|
- 找不到 `klayout`：安装系统级 KLayout 并加入 PATH；或使用回退（gdstk+SVG）。
|
||||||
|
- `cairosvg`/`gdstk` 报错：升级版本、确认写权限、检查输出目录存在。
|
||||||
|
- 训练集为空：检查 `paths.layout_dir` 与 `synthetic.png_dir` 是否存在且包含 .png；若 syn 目录为空将自动仅用真实数据。
|
||||||
|
|
||||||
|
---
|
||||||
|
|
||||||
|
## 🧪 合成数据管线与可视化
|
||||||
|
|
||||||
|
### 1) 生成合成 GDS
|
||||||
|
```bash
|
||||||
|
uv run python tools/generate_synthetic_layouts.py --out_dir data/synthetic/gds --num 200 --seed 42
|
||||||
|
```
|
||||||
|
|
||||||
|
### 2) 批量转换 GDS → PNG
|
||||||
|
```bash
|
||||||
|
uv run python tools/layout2png.py --in data/synthetic/gds --out data/synthetic/png --dpi 600
|
||||||
|
```
|
||||||
|
|
||||||
|
若本机未安装 KLayout，将自动回退到 gdstk+SVG 路径；图像外观可能与 KLayout 有差异。
|
||||||
|
|
||||||
|
### 3) 开启训练混采
|
||||||
|
在 `configs/base_config.yaml` 中设置：
|
||||||
|
```yaml
|
||||||
|
synthetic:
|
||||||
|
enabled: true
|
||||||
|
png_dir: data/synthetic/png
|
||||||
|
ratio: 0.3
|
||||||
|
```
|
||||||
|
|
||||||
|
### 4) 预览训练对（目检增强/H 一致性）
|
||||||
|
```bash
|
||||||
|
uv run python tools/preview_dataset.py --dir data/synthetic/png --out preview.png --n 8 --elastic
|
||||||
|
```
|
||||||
|
|
||||||
|
### 5) 开启/调整 Elastic 变形
|
||||||
|
```yaml
|
||||||
|
augment:
|
||||||
|
elastic:
|
||||||
|
enabled: true
|
||||||
|
alpha: 40
|
||||||
|
sigma: 6
|
||||||
|
alpha_affine: 6
|
||||||
|
prob: 0.3
|
||||||
|
photometric:
|
||||||
|
brightness_contrast: true
|
||||||
|
gauss_noise: true
|
||||||
|
```
|
||||||
|
|||||||
92
benchmark_grid.json
Normal file
92
benchmark_grid.json
Normal file
@@ -0,0 +1,92 @@
|
|||||||
|
[
|
||||||
|
{
|
||||||
|
"backbone": "vgg16",
|
||||||
|
"attention": "none",
|
||||||
|
"places": "backbone_high,desc_head",
|
||||||
|
"single_ms_mean": 351.6519069671631,
|
||||||
|
"single_ms_std": 1.8778125281542124,
|
||||||
|
"fpn_ms_mean": 719.3304697672526,
|
||||||
|
"fpn_ms_std": 3.949980966745213,
|
||||||
|
"runs": 3
|
||||||
|
},
|
||||||
|
{
|
||||||
|
"backbone": "vgg16",
|
||||||
|
"attention": "se",
|
||||||
|
"places": "backbone_high,desc_head",
|
||||||
|
"single_ms_mean": 349.7585455576579,
|
||||||
|
"single_ms_std": 1.9950684383137551,
|
||||||
|
"fpn_ms_mean": 721.4130560557047,
|
||||||
|
"fpn_ms_std": 2.7448351792281374,
|
||||||
|
"runs": 3
|
||||||
|
},
|
||||||
|
{
|
||||||
|
"backbone": "vgg16",
|
||||||
|
"attention": "cbam",
|
||||||
|
"places": "backbone_high,desc_head",
|
||||||
|
"single_ms_mean": 354.4490337371826,
|
||||||
|
"single_ms_std": 1.4903953036396786,
|
||||||
|
"fpn_ms_mean": 744.7629769643148,
|
||||||
|
"fpn_ms_std": 29.3233387791729,
|
||||||
|
"runs": 3
|
||||||
|
},
|
||||||
|
{
|
||||||
|
"backbone": "resnet34",
|
||||||
|
"attention": "none",
|
||||||
|
"places": "backbone_high,desc_head",
|
||||||
|
"single_ms_mean": 90.98696708679199,
|
||||||
|
"single_ms_std": 0.41179110533866975,
|
||||||
|
"fpn_ms_mean": 117.2173023223877,
|
||||||
|
"fpn_ms_std": 0.40632490569423124,
|
||||||
|
"runs": 3
|
||||||
|
},
|
||||||
|
{
|
||||||
|
"backbone": "resnet34",
|
||||||
|
"attention": "se",
|
||||||
|
"places": "backbone_high,desc_head",
|
||||||
|
"single_ms_mean": 90.78375498453777,
|
||||||
|
"single_ms_std": 0.4705899743190883,
|
||||||
|
"fpn_ms_mean": 115.90576171875,
|
||||||
|
"fpn_ms_std": 1.3081578935341862,
|
||||||
|
"runs": 3
|
||||||
|
},
|
||||||
|
{
|
||||||
|
"backbone": "resnet34",
|
||||||
|
"attention": "cbam",
|
||||||
|
"places": "backbone_high,desc_head",
|
||||||
|
"single_ms_mean": 96.49538993835449,
|
||||||
|
"single_ms_std": 3.17170034860506,
|
||||||
|
"fpn_ms_mean": 111.08938852945964,
|
||||||
|
"fpn_ms_std": 1.0126843546619573,
|
||||||
|
"runs": 3
|
||||||
|
},
|
||||||
|
{
|
||||||
|
"backbone": "efficientnet_b0",
|
||||||
|
"attention": "none",
|
||||||
|
"places": "backbone_high,desc_head",
|
||||||
|
"single_ms_mean": 40.451606114705406,
|
||||||
|
"single_ms_std": 1.5293525027201111,
|
||||||
|
"fpn_ms_mean": 127.30161348978679,
|
||||||
|
"fpn_ms_std": 0.08508800981401025,
|
||||||
|
"runs": 3
|
||||||
|
},
|
||||||
|
{
|
||||||
|
"backbone": "efficientnet_b0",
|
||||||
|
"attention": "se",
|
||||||
|
"places": "backbone_high,desc_head",
|
||||||
|
"single_ms_mean": 46.480417251586914,
|
||||||
|
"single_ms_std": 0.2622188910897682,
|
||||||
|
"fpn_ms_mean": 142.35099156697592,
|
||||||
|
"fpn_ms_std": 6.611047958580852,
|
||||||
|
"runs": 3
|
||||||
|
},
|
||||||
|
{
|
||||||
|
"backbone": "efficientnet_b0",
|
||||||
|
"attention": "cbam",
|
||||||
|
"places": "backbone_high,desc_head",
|
||||||
|
"single_ms_mean": 47.10610707600912,
|
||||||
|
"single_ms_std": 0.47150733957171853,
|
||||||
|
"fpn_ms_mean": 150.99199612935385,
|
||||||
|
"fpn_ms_std": 12.465987661773038,
|
||||||
|
"runs": 3
|
||||||
|
}
|
||||||
|
]
|
||||||
@@ -51,3 +51,24 @@ paths:
|
|||||||
val_ann_dir: "path/to/val/annotations"
|
val_ann_dir: "path/to/val/annotations"
|
||||||
template_dir: "path/to/templates"
|
template_dir: "path/to/templates"
|
||||||
model_path: "path/to/save/model_final.pth"
|
model_path: "path/to/save/model_final.pth"
|
||||||
|
|
||||||
|
# 数据增强与合成数据配置（可选）
|
||||||
|
augment:
|
||||||
|
elastic:
|
||||||
|
enabled: false
|
||||||
|
alpha: 40
|
||||||
|
sigma: 6
|
||||||
|
alpha_affine: 6
|
||||||
|
prob: 0.3
|
||||||
|
photometric:
|
||||||
|
brightness_contrast: true
|
||||||
|
gauss_noise: true
|
||||||
|
|
||||||
|
synthetic:
|
||||||
|
enabled: false
|
||||||
|
png_dir: "data/synthetic/png"
|
||||||
|
ratio: 0.0 # 0~1，训练时混合的合成样本比例
|
||||||
|
diffusion:
|
||||||
|
enabled: false
|
||||||
|
png_dir: "data/synthetic_diff/png"
|
||||||
|
ratio: 0.0 # 0~1，训练时混合的扩散样本比例
|
||||||
|
|||||||
@@ -1,6 +1,6 @@
|
|||||||
import os
|
import os
|
||||||
import json
|
import json
|
||||||
from typing import Tuple
|
from typing import Tuple, Optional
|
||||||
|
|
||||||
import cv2
|
import cv2
|
||||||
import numpy as np
|
import numpy as np
|
||||||
@@ -70,6 +70,8 @@ class ICLayoutTrainingDataset(Dataset):
|
|||||||
patch_size: int = 256,
|
patch_size: int = 256,
|
||||||
transform=None,
|
transform=None,
|
||||||
scale_range: Tuple[float, float] = (1.0, 1.0),
|
scale_range: Tuple[float, float] = (1.0, 1.0),
|
||||||
|
use_albu: bool = False,
|
||||||
|
albu_params: Optional[dict] = None,
|
||||||
) -> None:
|
) -> None:
|
||||||
self.image_dir = image_dir
|
self.image_dir = image_dir
|
||||||
self.image_paths = [
|
self.image_paths = [
|
||||||
@@ -80,6 +82,28 @@ class ICLayoutTrainingDataset(Dataset):
|
|||||||
self.patch_size = patch_size
|
self.patch_size = patch_size
|
||||||
self.transform = transform
|
self.transform = transform
|
||||||
self.scale_range = scale_range
|
self.scale_range = scale_range
|
||||||
|
# 可选的 albumentations 管道
|
||||||
|
self.albu = None
|
||||||
|
if use_albu:
|
||||||
|
try:
|
||||||
|
import albumentations as A # 延迟导入，避免环境未安装时报错
|
||||||
|
p = albu_params or {}
|
||||||
|
elastic_prob = float(p.get("prob", 0.3))
|
||||||
|
alpha = float(p.get("alpha", 40))
|
||||||
|
sigma = float(p.get("sigma", 6))
|
||||||
|
alpha_affine = float(p.get("alpha_affine", 6))
|
||||||
|
use_bc = bool(p.get("brightness_contrast", True))
|
||||||
|
use_noise = bool(p.get("gauss_noise", True))
|
||||||
|
transforms_list = [
|
||||||
|
A.ElasticTransform(alpha=alpha, sigma=sigma, alpha_affine=alpha_affine, p=elastic_prob),
|
||||||
|
]
|
||||||
|
if use_bc:
|
||||||
|
transforms_list.append(A.RandomBrightnessContrast(p=0.5))
|
||||||
|
if use_noise:
|
||||||
|
transforms_list.append(A.GaussNoise(var_limit=(5.0, 20.0), p=0.3))
|
||||||
|
self.albu = A.Compose(transforms_list)
|
||||||
|
except Exception:
|
||||||
|
self.albu = None
|
||||||
|
|
||||||
def __len__(self) -> int:
|
def __len__(self) -> int:
|
||||||
return len(self.image_paths)
|
return len(self.image_paths)
|
||||||
@@ -102,7 +126,13 @@ class ICLayoutTrainingDataset(Dataset):
|
|||||||
patch = image.crop((x, y, x + crop_size, y + crop_size))
|
patch = image.crop((x, y, x + crop_size, y + crop_size))
|
||||||
patch = patch.resize((self.patch_size, self.patch_size), Image.Resampling.LANCZOS)
|
patch = patch.resize((self.patch_size, self.patch_size), Image.Resampling.LANCZOS)
|
||||||
|
|
||||||
# 亮度/对比度增强
|
# photometric/elastic（在几何 H 之前）
|
||||||
|
patch_np_uint8 = np.array(patch)
|
||||||
|
if self.albu is not None:
|
||||||
|
patch_np_uint8 = self.albu(image=patch_np_uint8)["image"]
|
||||||
|
patch = Image.fromarray(patch_np_uint8)
|
||||||
|
else:
|
||||||
|
# 原有轻量光度增强
|
||||||
if np.random.random() < 0.5:
|
if np.random.random() < 0.5:
|
||||||
brightness_factor = np.random.uniform(0.8, 1.2)
|
brightness_factor = np.random.uniform(0.8, 1.2)
|
||||||
patch = patch.point(lambda px: int(np.clip(px * brightness_factor, 0, 255)))
|
patch = patch.point(lambda px: int(np.clip(px * brightness_factor, 0, 255)))
|
||||||
@@ -116,7 +146,6 @@ class ICLayoutTrainingDataset(Dataset):
|
|||||||
noise = np.random.normal(0, 5, patch_np.shape)
|
noise = np.random.normal(0, 5, patch_np.shape)
|
||||||
patch_np = np.clip(patch_np + noise, 0, 255)
|
patch_np = np.clip(patch_np + noise, 0, 255)
|
||||||
patch = Image.fromarray(patch_np.astype(np.uint8))
|
patch = Image.fromarray(patch_np.astype(np.uint8))
|
||||||
|
|
||||||
patch_np_uint8 = np.array(patch)
|
patch_np_uint8 = np.array(patch)
|
||||||
|
|
||||||
# 随机旋转与镜像（8个离散变换）
|
# 随机旋转与镜像（8个离散变换）
|
||||||
|
|||||||
317
docs/NextStep.md
317
docs/NextStep.md
@@ -1,173 +1,200 @@
|
|||||||
# 下一步工作计划 (NextStep)
|
## 一、数据策略与增强 (Data Strategy & Augmentation)
|
||||||
|
|
||||||
**最后更新**: 2025-10-20
|
> 目标：提升模型的鲁棒性和泛化能力，减少对大量真实数据的依赖。
|
||||||
**范围**: 仅聚焦于 `feature_work.md` 的第二部分「模型架构 (Model Architecture)」的落地执行计划
|
|
||||||
**上下文**: 核心功能已完成，本文档将模型架构优化转化为可执行的工程计划，便于直接实施与验收。
|
|
||||||
|
|
||||||
> 参考来源：`docs/feature_work.md` 第二部分；更宏观的阶段规划见 `docs/todos/`
|
- [x] 引入弹性变形 (Elastic Transformations)
|
||||||
|
- ✔️ 价值：模拟芯片制造中可能出现的微小物理形变，使模型对非刚性变化更鲁棒。
|
||||||
|
- 🧭 关键原则（与当前数据管线一致）：
|
||||||
|
- 现有自监督训练数据集 `ICLayoutTrainingDataset` 会返回 (original, rotated, H)；其中 H 是两张 patch 间的单应关系，用于 loss 监督。
|
||||||
|
- 非刚性弹性变形若只对其中一张或在生成 H 之后施加，会破坏几何约束，导致 H 失效。
|
||||||
|
- 因此，Elastic 需在“生成 homography 配对之前”对基础 patch 施加；随后对该已变形的 patch 再执行旋转/镜像与单应计算，这样 H 仍严格成立。
|
||||||
|
- 📝 执行计划：
|
||||||
|
1) 依赖核对
|
||||||
|
- `pyproject.toml` 已包含 `albumentations>=2.0.8`，无需新增依赖；确保环境安装齐全。
|
||||||
|
2) 集成位置与方式
|
||||||
|
- 在 `data/ic_dataset.py` 的 `ICLayoutTrainingDataset.__getitem__` 中，裁剪并缩放得到 `patch` 后，转换为 `np.ndarray`，对其调用 `albumentations` 管道（包含 `A.ElasticTransform`）。
|
||||||
|
- 将变形后的 `patch_np_uint8` 作为“基准图”，再按现有逻辑计算旋转/镜像与 `homography`，生成 `transformed_patch`，从而确保 H 有效。
|
||||||
|
3) 代码改动清单（建议）
|
||||||
|
- `data/ic_dataset.py`
|
||||||
|
- 顶部新增：`import albumentations as A`
|
||||||
|
- `__init__` 新增可选参数：`use_albu: bool=False`、`albu_params: dict|None=None`
|
||||||
|
- 在 `__init__` 构造 `self.albu = A.Compose([...])`（当 `use_albu` 为 True 时），包含：
|
||||||
|
- `A.ElasticTransform(alpha=40, sigma=6, alpha_affine=6, p=0.3)`
|
||||||
|
- （可选）`A.RandomBrightnessContrast(p=0.5)`、`A.GaussNoise(var_limit=(5.0, 20.0), p=0.3)` 以替代当前手写的亮度/对比度与噪声逻辑（减少重复）。
|
||||||
|
- 在 `__getitem__`：裁剪与缩放后，若启用 `self.albu`：`patch_np_uint8 = self.albu(image=patch_np_uint8)["image"]`，随后再计算旋转/镜像与 `homography`。
|
||||||
|
- 注意：保持输出张量与当前 `utils.data_utils.get_transform()` 兼容（单通道→三通道→Normalize）。
|
||||||
|
- `configs/base_config.yaml`
|
||||||
|
- 新增配置段：
|
||||||
|
- `augment.elastic.enabled: true|false`
|
||||||
|
- `augment.elastic.alpha: 40`
|
||||||
|
- `augment.elastic.sigma: 6`
|
||||||
|
- `augment.elastic.alpha_affine: 6`
|
||||||
|
- `augment.elastic.prob: 0.3`
|
||||||
|
- （可选）`augment.photometric.*` 开关与参数
|
||||||
|
- `train.py`
|
||||||
|
- 从配置读取上述参数，并将 `use_albu` 与 `albu_params` 通过 `ICLayoutTrainingDataset(...)` 传入（不影响现有 `get_transform()`）。
|
||||||
|
4) 参数与默认值建议
|
||||||
|
- 起始：`alpha=40, sigma=6, alpha_affine=6, p=0.3`；根据训练收敛与可视化效果微调。
|
||||||
|
- 若发现描述子对局部形变敏感，可逐步提高 `alpha` 或 `p`；若训练不稳定则降低。
|
||||||
|
5) 验证与可视化
|
||||||
|
- 在 `tests/benchmark_grid.py` 或新增简单可视化脚本中，采样 16 个 (original, rotated) 对，叠加可视化 H 变换后的网格，确认几何一致性未破坏。
|
||||||
|
- 训练前 1000 个 batch：记录 `loss_det/loss_desc` 曲线，确认未出现异常发散。
|
||||||
|
|
||||||
---
|
- [x] 创建合成版图数据生成器
|
||||||
|
- ✔️ 价值：解决真实版图数据获取难、数量少的问题，通过程序化生成大量多样化的训练样本。
|
||||||
|
- 📝 执行计划：
|
||||||
|
1) 新增脚本 `tools/generate_synthetic_layouts.py`
|
||||||
|
- 目标：使用 `gdstk` 程序化生成包含不同尺寸、密度与单元类型的 GDSII 文件。
|
||||||
|
- 主要能力：
|
||||||
|
- 随机生成“标准单元”模版（如若干矩形/多边形组合）、金属走线、过孔阵列；
|
||||||
|
- 支持多层（layer/datatype）与规则化阵列（row/col pitch）、占空比（density）控制；
|
||||||
|
- 形状参数与布局由随机种子控制，支持可重复性。
|
||||||
|
- CLI 设计（示例）：
|
||||||
|
- `--out-dir data/synthetic/gds`、`--num-samples 1000`、`--seed 42`
|
||||||
|
- 版图规格：`--width 200um --height 200um --grid 0.1um`
|
||||||
|
- 多样性开关：`--cell-types NAND,NOR,INV --metal-layers 3 --density 0.1-0.6`
|
||||||
|
- 关键实现要点：
|
||||||
|
- 使用 `gdstk.Library()` 与 `gdstk.Cell()` 组装基本单元；
|
||||||
|
- 通过 `gdstk.Reference` 和阵列生成放置；
|
||||||
|
- 生成完成后 `library.write_gds(path)` 落盘。
|
||||||
|
2) 批量转换 GDSII → PNG（训练用）
|
||||||
|
- 现状核对：仓库中暂无 `tools/layout2png.py`；计划新增该脚本（与本项一并交付）。
|
||||||
|
- 推荐实现 A（首选）：使用 `klayout` 的 Python API（`pya`）以无头模式加载 GDS，指定层映射与缩放，导出为高分辨率 PNG：
|
||||||
|
- 脚本 `tools/layout2png.py` 提供 CLI：`--in data/synthetic/gds --out data/synthetic/png --dpi 600 --layers 1/0:gray,2/0:blue ...`
|
||||||
|
- 支持目录批量与单文件转换；可配置画布背景、线宽、边距。
|
||||||
|
- 替代实现 B：导出 SVG 再用 `cairosvg` 转 PNG（依赖已在项目中），适合无 klayout 环境的场景。
|
||||||
|
- 输出命名规范：与 GDS 同名，如 `chip_000123.gds → chip_000123.png`。
|
||||||
|
3) 数据目录与元数据
|
||||||
|
- 目录结构建议：
|
||||||
|
- `data/synthetic/gds/`、`data/synthetic/png/`、`data/synthetic/meta/`
|
||||||
|
- 可选：为每个样本生成 `meta/*.json`，记录层数、单元类型分布、密度等，用于后续分析/分层采样。
|
||||||
|
4) 与训练集集成
|
||||||
|
- `configs/base_config.yaml` 新增：
|
||||||
|
- `paths.synthetic_dir: data/synthetic/png`
|
||||||
|
- `training.use_synthetic_ratio: 0.0~1.0`（混合采样比例；例如 0.3 表示 30% 合成样本）
|
||||||
|
- 在 `train.py` 中：
|
||||||
|
- 若 `use_synthetic_ratio>0`，构建一个 `ICLayoutTrainingDataset` 指向合成 PNG 目录；
|
||||||
|
- 实现简单的比例采样器或 `ConcatDataset + WeightedRandomSampler` 以按比例混合真实与合成样本。
|
||||||
|
5) 质量与稳健性检查
|
||||||
|
- 可视化抽样：随机展示若干 PNG，检查层次颜色、对比度、线宽是否清晰；
|
||||||
|
- 分布对齐：统计真实数据与合成数据的连线长度分布、拓扑度量（如节点度、环路数量），做基础分布对齐；
|
||||||
|
- 训练烟雾测试：仅用 100～200 个合成样本跑 1～2 个 epoch，确认训练闭环无错误、loss 正常下降。
|
||||||
|
6) 基准验证与复盘
|
||||||
|
- 在 `tests/benchmark_grid.py` 与 `tests/benchmark_backbones.py` 增加一组“仅真实 / 真实+合成”的对照实验；
|
||||||
|
- 记录 mAP/匹配召回/描述子一致性等指标，评估增益；
|
||||||
|
- 产出 `docs/Performance_Benchmark.md` 的对比表格。
|
||||||
|
|
||||||
## 🔴 模型架构优化（Feature Work 第二部分）
|
### 验收标准 (Acceptance Criteria)
|
||||||
|
|
||||||
目标：在保证现有精度的前提下，提升特征提取效率与推理速度；为后续注意力机制与多尺度策略提供更强的特征基础。
|
- Elastic 变形：
|
||||||
|
- [ ] 训练数据可视化（含 H 网格叠加）无几何错位；
|
||||||
|
- [ ] 训练前若干 step loss 无异常尖峰，长期收敛不劣于 baseline；
|
||||||
|
- [ ] 可通过配置无缝开/关与调参。
|
||||||
|
- 合成数据：
|
||||||
|
- [ ] 能批量生成带多层元素的 GDS 文件并成功转为 PNG；
|
||||||
|
- [ ] 训练脚本可按设定比例混合采样真实与合成样本；
|
||||||
|
- [ ] 在小规模对照实验中，验证指标有稳定或可解释的变化（不劣化）。
|
||||||
|
|
||||||
### 总体验收标准（全局）
|
### 风险与规避 (Risks & Mitigations)
|
||||||
- [ ] 训练/验证流程在新骨干和注意力方案下均可跑通，无崩溃/NaN。
|
|
||||||
- [ ] 在代表性验证集上，最终指标（IoU/mAP）不低于当前 VGG-16 基线；若下降需给出改进措施或回滚建议。
|
|
||||||
- [ ] 推理时延或显存占用至少一种维度优于基线，或达到“相当 + 结构可扩展”的工程收益。
|
|
||||||
- [ ] 关键改动均通过配置开关控制，可随时回退。
|
|
||||||
|
|
||||||
---
|
- 非刚性变形破坏 H 的风险：仅在生成 homography 前对基准 patch 施加 Elastic，或在两图上施加相同变形但更新 H′=f∘H∘f⁻¹（当前计划采用前者，简单且稳定）。
|
||||||
|
- GDS → PNG 渲染差异：优先使用 `klayout`，保持工业级渲染一致性；无 `klayout` 时使用 SVG→PNG 备选路径。
|
||||||
|
- 合成分布与真实分布不匹配：通过密度与单元类型分布约束进行对齐，并在训练中控制混合比例渐进提升。
|
||||||
|
|
||||||
## 2.1 实验更现代的骨干网络（Backbone）
|
### 里程碑与时间估算 (Milestones & ETA)
|
||||||
|
|
||||||
优先级：🟠 中 | 预计工期：~1 周 | 产出：可切换的 backbone 实现 + 对照报告
|
## 二、实现状态与使用说明（2025-10-20 更新）
|
||||||
|
|
||||||
### 设计要点（小合约）
|
- Elastic 变形已按计划集成：
|
||||||
- 输入：与现有 `RoRD` 一致的图像张量 B×C×H×W。
|
- 开关与参数：见 `configs/base_config.yaml` 下的 `augment.elastic` 与 `augment.photometric`；
|
||||||
- 输出：供检测头/描述子头使用的中高层特征张量；通道数因骨干不同而异（VGG:512、ResNet34:512、Eff-B0:1280）。
|
- 数据集实现：`data/ic_dataset.py` 中 `ICLayoutTrainingDataset`；
|
||||||
- 约束：不改变下游头部的接口形状（头部输入通道需根据骨干进行对齐适配）。
|
- 可视化验证：`tools/preview_dataset.py --dir <png_dir> --n 8 --elastic`。
|
||||||
- 失败模式：通道不匹配/梯度不通/预训练权重未正确加载/收敛缓慢。
|
|
||||||
|
|
||||||
### 配置扩展（YAML）
|
- 合成数据生成与渲染：
|
||||||
在 `configs/base_config.yaml` 增加（或确认存在）：
|
- 生成 GDS：`tools/generate_synthetic_layouts.py --out-dir data/synthetic/gds --num 100 --seed 42`；
|
||||||
|
- 转换 PNG：`tools/layout2png.py --in data/synthetic/gds --out data/synthetic/png --dpi 600`；
|
||||||
|
- 训练混采：在 `configs/base_config.yaml` 设置 `synthetic.enabled: true`、`synthetic.png_dir: data/synthetic/png`、`synthetic.ratio: 0.3`。
|
||||||
|
|
||||||
```yaml
|
- 训练脚本：
|
||||||
model:
|
- `train.py` 已接入真实/合成混采（ConcatDataset + WeightedRandomSampler），验证集仅用真实数据；
|
||||||
backbone:
|
- TensorBoard 文本摘要记录数据构成（mix 开关、比例、样本量）。
|
||||||
name: "vgg16" # 可选：vgg16 | resnet34 | efficientnet_b0
|
|
||||||
pretrained: true
|
注意：若未安装 KLayout，可自动回退 gdstk+SVG 路径；显示效果可能与 KLayout 存在差异。
|
||||||
# 用于选择抽取的特征层（按不同骨干约定名称）
|
|
||||||
feature_layers:
|
- D1：Elastic 集成 + 可视化验证（代码改动与测试）
|
||||||
vgg16: ["relu3_3", "relu4_3"]
|
- D2：合成生成器初版（GDS 生成 + PNG 渲染脚本）
|
||||||
resnet34: ["layer3", "layer4"]
|
- D3：训练混合采样接入 + 小规模基准
|
||||||
efficientnet_b0: ["features_5", "features_7"]
|
- D4：参数扫与报告更新（Performance_Benchmark.md）
|
||||||
|
|
||||||
|
### 一键流水线（生成 → 渲染 → 预览 → 训练）
|
||||||
|
|
||||||
|
1) 生成 GDS（合成版图）
|
||||||
|
```bash
|
||||||
|
uv run python tools/generate_synthetic_layouts.py --out_dir data/synthetic/gds --num 200 --seed 42
|
||||||
```
|
```
|
||||||
|
|
||||||
### 代码改动建议
|
2) 渲染 PNG（KLayout 优先，自动回退 gdstk+SVG）
|
||||||
- 文件：`models/rord.py`
|
```bash
|
||||||
1) 在 `__init__` 中根据 `cfg.model.backbone.name` 动态构建骨干：
|
uv run python tools/layout2png.py --in data/synthetic/gds --out data/synthetic/png --dpi 600
|
||||||
- vgg16（现状保持）
|
```
|
||||||
- resnet34：从 `torchvision.models.resnet34(weights=IMAGENET1K_V1)` 构建；保存 `layer3/layer4` 输出。
|
|
||||||
- efficientnet_b0：从 `torchvision.models.efficientnet_b0(weights=IMAGENET1K_V1)` 构建；保存末两段 `features` 输出。
|
|
||||||
2) 为不同骨干提供统一的“中间层特征导出”接口（注册 forward hook 或显式调用子模块）。
|
|
||||||
3) 依据所选骨干的输出通道，调整检测头与描述子头的输入通道（如使用 1×1 conv 过渡层以解耦通道差异）。
|
|
||||||
4) 保持现有前向签名与返回数据结构不变（训练/推理兼容）。
|
|
||||||
|
|
||||||
### 进展更新（2025-10-20）
|
3) 预览训练对（核验增强/H 一致性）
|
||||||
- 已完成：在 `models/rord.py` 集成多骨干选择（`vgg16`/`resnet34`/`efficientnet_b0`），并实现统一的中间层抽取函数 `_extract_c234`（可后续重构为 `build_backbone`/`extract_features` 明确接口）。
|
```bash
|
||||||
- 已完成：FPN 通用化，基于 C2/C3/C4 构建 P2/P3/P4，按骨干返回正确的 stride。
|
uv run python tools/preview_dataset.py --dir data/synthetic/png --out preview.png --n 8 --elastic
|
||||||
- 已完成：单图前向 Smoke Test（三种骨干，单尺度与 FPN）均通过。
|
```
|
||||||
- 已完成：CPU 环境 A/B 基准（单尺度 vs FPN）见 `docs/description/Performance_Benchmark.md`。
|
|
||||||
- 待完成：GPU 环境基准（速度/显存）、基于真实数据的精度评估与收敛曲线对比。
|
|
||||||
|
|
||||||
### 落地步骤（Checklist）
|
4) 在 YAML 中开启混采与 Elastic（示例）
|
||||||
- [x] 在 `models/rord.py` 增加/落地骨干构建与中间层抽取逻辑（当前通过 `_extract_c234` 实现）。
|
|
||||||
- [x] 接入 ResNet-34：返回等价中高层特征（layer2/3/4，通道≈128/256/512）。
|
|
||||||
- [x] 接入 EfficientNet-B0：返回 `features[2]/[3]/[6]`（约 24/40/192），FPN 以 1×1 横向连接对齐到 `fpn_out_channels`。
|
|
||||||
- [x] 头部适配：单尺度头使用骨干高层通道数；FPN 头统一使用 `fpn_out_channels`。
|
|
||||||
- [ ] 预训练权重：支持 `pretrained=true` 加载；补充权重加载摘要打印（哪些层未命中）。
|
|
||||||
- [x] 单图 smoke test：前向通过、无 NaN（三种骨干，单尺度与 FPN）。
|
|
||||||
|
|
||||||
### 评测与选择（A/B 实验）
|
|
||||||
- [ ] 在固定数据与超参下，比较 vgg16/resnet34/efficientnet_b0：
|
|
||||||
- 收敛速度（loss 曲线 0-5 epoch）
|
|
||||||
- 推理速度（ms / 2048×2048）与显存（GB）[CPU 初步结果已产出，GPU 待复测；见 `docs/description/Performance_Benchmark.md`]
|
|
||||||
- 验证集 IoU/mAP（真实数据集待跑）
|
|
||||||
- [ ] 形成表格与可视化图，给出选择结论与原因（CPU 版初稿已在报告中给出观察）。
|
|
||||||
- [ ] 若新骨干在任一关键指标明显受损，则暂缓替换，仅保留为可切换实验选项。
|
|
||||||
|
|
||||||
### 验收标准（2.1）
|
|
||||||
- [ ] 三种骨干方案均可训练与推理（当前仅验证推理，训练与收敛待验证）；
|
|
||||||
- [ ] 最终入选骨干在 IoU/mAP 不低于 VGG 的前提下，带来显著的速度/显存优势之一；
|
|
||||||
- [x] 切换完全配置化（无需改代码）。
|
|
||||||
|
|
||||||
### 风险与回滚（2.1）
|
|
||||||
- 通道不匹配导致维度错误 → 在进入头部前统一使用 1×1 conv 适配；
|
|
||||||
- 预训练权重与自定义层名不一致 → 显式映射并记录未加载层；
|
|
||||||
- 收敛变慢 → 暂时提高训练轮数、调学习率/BN 冻结策略；不达标即回滚 `backbone.name=vgg16`。
|
|
||||||
|
|
||||||
---
|
|
||||||
|
|
||||||
## 2.2 集成注意力机制（CBAM / SE-Net）
|
|
||||||
|
|
||||||
优先级：🟠 中 | 预计工期：~7–10 天 | 产出：注意力增强的 RoRD 变体 + 对照报告
|
|
||||||
|
|
||||||
### 模块选择与嵌入位置
|
|
||||||
- 方案 A：CBAM（通道注意 + 空间注意），插入至骨干高层与两类头部之前；
|
|
||||||
- 方案 B：SE-Net（通道注意），轻量但仅通道维，插入多个阶段以增强稳定性；
|
|
||||||
- 建议：先实现 CBAM，保留 SE 作为备选开关。
|
|
||||||
|
|
||||||
### 配置扩展（YAML）
|
|
||||||
```yaml
|
```yaml
|
||||||
model:
|
synthetic:
|
||||||
attention:
|
|
||||||
enabled: true
|
enabled: true
|
||||||
type: "cbam" # 可选：cbam | se | none
|
png_dir: data/synthetic/png
|
||||||
places: ["backbone_high", "det_head", "desc_head"]
|
ratio: 0.3
|
||||||
# 可选超参：reduction、kernel_size 等
|
|
||||||
reduction: 16
|
augment:
|
||||||
spatial_kernel: 7
|
elastic:
|
||||||
|
enabled: true
|
||||||
|
alpha: 40
|
||||||
|
sigma: 6
|
||||||
|
alpha_affine: 6
|
||||||
|
prob: 0.3
|
||||||
```
|
```
|
||||||
|
|
||||||
### 代码改动建议
|
5) 开始训练
|
||||||
- 文件：`models/rord.py`
|
```bash
|
||||||
1) 实现 `CBAM` 与 `SEBlock` 模块（或从可靠实现迁移），提供简洁 forward。
|
uv run python train.py --config configs/base_config.yaml
|
||||||
2) 在 `__init__` 中依据 `cfg.model.attention` 决定在何处插入：
|
```
|
||||||
- backbone 高层输出后（增强高层语义的判别性）；
|
|
||||||
- 检测头、描述子头输入前（分别强化不同任务所需特征）。
|
|
||||||
3) 注意保持张量尺寸不变；若引入残差结构，保证与原路径等价时可退化为恒等映射。
|
|
||||||
|
|
||||||
### 落地步骤（Checklist）
|
可选：使用单脚本一键执行（含配置写回）
|
||||||
- [ ] 实现 `CBAM`：通道注意（MLP/Avg+Max Pool）+ 空间注意（7×7 conv）。
|
```bash
|
||||||
- [ ] 实现 `SEBlock`：Squeeze（全局池化）+ Excitation（MLP, reduction）。
|
uv run python tools/synth_pipeline.py --out_root data/synthetic --num 200 --dpi 600 \
|
||||||
- [ ] 在 `RoRD` 中用配置化开关插拔注意力，默认关闭。
|
--config configs/base_config.yaml --ratio 0.3 --enable_elastic
|
||||||
- [ ] 在进入检测/描述子头前分别测试开启/关闭注意力的影响。
|
```
|
||||||
- [ ] 记录注意力图（可选）：导出中间注意图用于可视化对比。
|
|
||||||
|
|
||||||
### 训练与评估
|
### 参数建议与经验
|
||||||
- [ ] 以入选骨干为基线，分别开启 `cbam` 与 `se` 进行对照；
|
|
||||||
- [ ] 记录：训练损失、验证 IoU/mAP、推理时延/显存；
|
|
||||||
- [ ] 观察注意力图是否集中在关键几何（边角/交点/突变）；
|
|
||||||
- [ ] 若带来过拟合迹象（验证下降），尝试减弱注意力强度或减少插入位置。
|
|
||||||
|
|
||||||
### 验收标准（2.2）
|
- 渲染 DPI：600–900 通常足够，图形极细时可提高到 1200（注意磁盘与 IO）。
|
||||||
- [ ] 模型在开启注意力后稳定训练，无数值异常；
|
- 混采比例 synthetic.ratio：
|
||||||
- [ ] 指标不低于无注意力基线；若提升则量化收益；
|
- 数据少（<500 张）可取 0.3–0.5；
|
||||||
- [ ] 配置可一键关闭以回退。
|
- 数据中等（500–2000 张）建议 0.2–0.3；
|
||||||
|
- 数据多（>2000 张）建议 0.1–0.2 以免分布偏移。
|
||||||
|
- Elastic 强度：从 alpha=40, sigma=6 开始；若描述子对局部形变敏感，可小步上调 alpha 或 prob。
|
||||||
|
|
||||||
### 风险与回滚（2.2）
|
### 质量检查清单（建议在首次跑通后执行）
|
||||||
- 注意力导致过拟合或梯度不稳 → 降低 reduction、减少插入点、启用正则；
|
|
||||||
- 推理时延上升明显 → 对注意力路径进行轻量化（如仅通道注意或更小 kernel）。
|
|
||||||
|
|
||||||
---
|
- 预览拼图无明显几何错位（orig/rot 对应边界对齐合理）。
|
||||||
|
- 训练日志包含混采信息（real/syn 样本量、ratio、启停状态）。
|
||||||
|
- 若开启 Elastic，训练初期 loss 无异常尖峰，长期收敛不劣于 baseline。
|
||||||
|
- 渲染 PNG 与 GDS 在关键层上形态一致（优先使用 KLayout）。
|
||||||
|
|
||||||
## 工程与度量配套
|
### 常见问题与排查（FAQ）
|
||||||
|
|
||||||
### 实验记录（建议）
|
- klayout: command not found
|
||||||
- 在 TensorBoard 中新增：
|
- 方案A：安装系统级 KLayout 并确保可执行文件在 PATH；
|
||||||
- `arch/backbone_name`、`arch/attention_type`（Text/Scalar）；
|
- 方案B：暂用 gdstk+SVG 回退（外观可能略有差异）。
|
||||||
- `train/loss_total`、`eval/iou_metric`、`eval/map`；
|
- cairosvg 报错或 SVG 不生成
|
||||||
- 推理指标：`infer/ms_per_image`、`infer/vram_gb`。
|
- 升级 `cairosvg` 与 `gdstk`；确保磁盘有写入权限；检查 `.svg` 是否被安全软件拦截。
|
||||||
|
- gdstk 版本缺少 write_svg
|
||||||
### 对照报告模板（最小集）
|
- 尝试升级 gdstk；脚本已做 library 与 cell 双路径兼容，仍失败则优先使用 KLayout。
|
||||||
- 数据集与配置摘要（随机种子、批大小、学习率、图像尺寸）。
|
- 训练集为空或样本过少
|
||||||
- 三个骨干 + 注意力开关的结果表（速度/显存/IoU/mAP）。
|
- 检查 `paths.layout_dir` 与 `synthetic.png_dir` 是否存在且包含 .png；ratio>0 但 syn 目录为空会自动回退仅真实数据。
|
||||||
- 结论与落地选择（保留/关闭/待进一步实验）。
|
|
||||||
|
|
||||||
---
|
|
||||||
|
|
||||||
## 排期与里程碑（建议）
|
|
||||||
- M1（1 天）：骨干切换基础设施与通道适配层；单图 smoke 测试。
|
|
||||||
- M2（2–3 天）：ResNet34 与 EfficientNet-B0 接入与跑通；
|
|
||||||
- M3（1–2 天）：A/B 评测与结论；
|
|
||||||
- M4（3–4 天）：注意力模块接入、训练对照、报告输出。
|
|
||||||
|
|
||||||
---
|
|
||||||
|
|
||||||
## 相关参考
|
|
||||||
- 源文档：`docs/feature_work.md` 第二部分（模型架构）
|
|
||||||
- 阶段规划：`docs/todos/`
|
|
||||||
- 配置系统：`configs/base_config.yaml`
|
|
||||||
|
|
||||||
|
|||||||
@@ -17,6 +17,53 @@
|
|||||||
|
|
||||||
- 备注：本次测试在 CPU 上进行，`gpu_mem_mb` 始终为 0。
|
- 备注：本次测试在 CPU 上进行，`gpu_mem_mb` 始终为 0。
|
||||||
|
|
||||||
|
## 注意力 A/B（CPU，resnet34，512×512，runs=10，places=backbone_high+desc_head）
|
||||||
|
|
||||||
|
| Attention | Single Mean ± Std | FPN Mean ± Std |
|
||||||
|
|-----------|-------------------:|----------------:|
|
||||||
|
| none | 97.57 ± 0.55 | 124.57 ± 0.48 |
|
||||||
|
| se | 101.48 ± 2.13 | 123.12 ± 0.50 |
|
||||||
|
| cbam | 119.80 ± 2.38 | 123.11 ± 0.71 |
|
||||||
|
|
||||||
|
观察：
|
||||||
|
- 单尺度路径对注意力类型更敏感，CBAM 开销相对更高，SE 较轻；
|
||||||
|
- FPN 路径耗时在本次设置下差异很小（可能因注意力仅在 `backbone_high/desc_head`，且 FPN 头部计算占比较高）。
|
||||||
|
|
||||||
|
复现实验：
|
||||||
|
```zsh
|
||||||
|
PYTHONPATH=. uv run python tests/benchmark_attention.py \
|
||||||
|
--device cpu --image-size 512 --runs 10 \
|
||||||
|
--backbone resnet34 --places backbone_high desc_head
|
||||||
|
```
|
||||||
|
|
||||||
|
## 三维基准（Backbone × Attention × Single/FPN）
|
||||||
|
|
||||||
|
环境：CPU，输入 1×3×512×512，重复 3 次，places=backbone_high,desc_head。
|
||||||
|
|
||||||
|
| Backbone | Attention | Single Mean ± Std (ms) | FPN Mean ± Std (ms) |
|
||||||
|
|------------------|-----------|-----------------------:|--------------------:|
|
||||||
|
| vgg16 | none | 351.65 ± 1.88 | 719.33 ± 3.95 |
|
||||||
|
| vgg16 | se | 349.76 ± 2.00 | 721.41 ± 2.74 |
|
||||||
|
| vgg16 | cbam | 354.45 ± 1.49 | 744.76 ± 29.32 |
|
||||||
|
| resnet34 | none | 90.99 ± 0.41 | 117.22 ± 0.41 |
|
||||||
|
| resnet34 | se | 90.78 ± 0.47 | 115.91 ± 1.31 |
|
||||||
|
| resnet34 | cbam | 96.50 ± 3.17 | 111.09 ± 1.01 |
|
||||||
|
| efficientnet_b0 | none | 40.45 ± 1.53 | 127.30 ± 0.09 |
|
||||||
|
| efficientnet_b0 | se | 46.48 ± 0.26 | 142.35 ± 6.61 |
|
||||||
|
| efficientnet_b0 | cbam | 47.11 ± 0.47 | 150.99 ± 12.47 |
|
||||||
|
|
||||||
|
复现实验：
|
||||||
|
|
||||||
|
```zsh
|
||||||
|
PYTHONPATH=. uv run python tests/benchmark_grid.py \
|
||||||
|
--device cpu --image-size 512 --runs 3 \
|
||||||
|
--backbones vgg16 resnet34 efficientnet_b0 \
|
||||||
|
--attentions none se cbam \
|
||||||
|
--places backbone_high desc_head
|
||||||
|
```
|
||||||
|
|
||||||
|
运行会同时输出控制台摘要并保存 JSON：`benchmark_grid.json`。
|
||||||
|
|
||||||
## 观察与解读
|
## 观察与解读
|
||||||
- vgg16 明显最慢，FPN 额外的横向/上采样代价在 CPU 上更突出（>2×）。
|
- vgg16 明显最慢，FPN 额外的横向/上采样代价在 CPU 上更突出（>2×）。
|
||||||
- resnet34 在单尺度上显著快于 vgg16，FPN 增幅较小（约 +25%）。
|
- resnet34 在单尺度上显著快于 vgg16，FPN 增幅较小（约 +25%）。
|
||||||
|
|||||||
@@ -1,5 +1,41 @@
|
|||||||
# 后续工作
|
# 后续工作
|
||||||
|
|
||||||
|
## 新增功能汇总（2025-10-20）
|
||||||
|
|
||||||
|
- 数据增强：集成 `albumentations` 的 ElasticTransform（配置在 `augment.elastic`），并保持几何配对的 H 正确性。
|
||||||
|
- 合成数据：新增 `tools/generate_synthetic_layouts.py`（GDS 生成）与 `tools/layout2png.py`（GDS→PNG 批量转换）。
|
||||||
|
- 训练混采：`train.py` 接入真实/合成混采，按 `synthetic.ratio` 使用加权采样；验证集仅使用真实数据。
|
||||||
|
- 可视化：`tools/preview_dataset.py` 快速导出训练对的拼图图，便于人工质检。
|
||||||
|
|
||||||
|
## 立即可做的小改进
|
||||||
|
|
||||||
|
- 在 `layout2png.py` 增加图层配色与线宽配置（读取 layermap 或命令行参数）。
|
||||||
|
- 为 `ICLayoutTrainingDataset` 添加随机裁剪失败时的回退逻辑（极小图像）。
|
||||||
|
- 增加最小单元测试：验证 ElasticTransform 下 H 的 warp 一致性（采样角点/网格点）。
|
||||||
|
- 在 README 增加一键命令合集（生成合成数据 → 渲染 → 预览 → 训练）。
|
||||||
|
|
||||||
|
## 一键流程与排查（摘要）
|
||||||
|
|
||||||
|
**一键命令**：
|
||||||
|
```bash
|
||||||
|
uv run python tools/generate_synthetic_layouts.py --out_dir data/synthetic/gds --num 200 --seed 42
|
||||||
|
uv run python tools/layout2png.py --in data/synthetic/gds --out data/synthetic/png --dpi 600
|
||||||
|
uv run python tools/preview_dataset.py --dir data/synthetic/png --out preview.png --n 8 --elastic
|
||||||
|
uv run python train.py --config configs/base_config.yaml
|
||||||
|
```
|
||||||
|
|
||||||
|
或使用单脚本一键执行（含配置写回）：
|
||||||
|
```bash
|
||||||
|
uv run python tools/synth_pipeline.py --out_root data/synthetic --num 200 --dpi 600 \
|
||||||
|
--config configs/base_config.yaml --ratio 0.3 --enable_elastic
|
||||||
|
```
|
||||||
|
|
||||||
|
**参数建议**：DPI=600–900；ratio=0.2–0.3（首训）；Elastic 从 alpha=40/sigma=6 起步。
|
||||||
|
|
||||||
|
**FAQ**：
|
||||||
|
- 找不到 klayout：安装后确保在 PATH；无则使用回退渲染（外观可能有差异）。
|
||||||
|
- SVG/PNG 未生成：检查写权限与版本（cairosvg/gdstk），或优先用 KLayout。
|
||||||
|
|
||||||
本文档整合了 RoRD 项目的优化待办清单和训练需求，用于规划未来的开发和实验工作。
|
本文档整合了 RoRD 项目的优化待办清单和训练需求，用于规划未来的开发和实验工作。
|
||||||
|
|
||||||
---
|
---
|
||||||
@@ -12,18 +48,70 @@
|
|||||||
|
|
||||||
> *目标：提升模型的鲁棒性和泛化能力，减少对大量真实数据的依赖。*
|
> *目标：提升模型的鲁棒性和泛化能力，减少对大量真实数据的依赖。*
|
||||||
|
|
||||||
- [ ] **引入弹性变形 (Elastic Transformations)**
|
- [x] **引入弹性变形 (Elastic Transformations)**
|
||||||
- **✔️ 价值**: 模拟芯片制造中可能出现的微小物理形变，使模型对非刚性变化更鲁棒。
|
- **✔️ 价值**: 模拟芯片制造中可能出现的微小物理形变，使模型对非刚性变化更鲁棒。
|
||||||
- **📝 执行方案**:
|
- **📝 执行方案**:
|
||||||
1. 添加 `albumentations` 库作为项目依赖。
|
1. 添加 `albumentations` 库作为项目依赖。
|
||||||
2. 在 `train.py` 的 `ICLayoutTrainingDataset` 类中，集成 `A.ElasticTransform` 到数据增强管道中。
|
2. 在 `train.py` 的 `ICLayoutTrainingDataset` 类中，集成 `A.ElasticTransform` 到数据增强管道中。
|
||||||
- [ ] **创建合成版图数据生成器**
|
- [x] **创建合成版图数据生成器**
|
||||||
- **✔️ 价值**: 解决真实版图数据获取难、数量少的问题，通过程序化生成大量多样化的训练样本。
|
- **✔️ 价值**: 解决真实版图数据获取难、数量少的问题，通过程序化生成大量多样化的训练样本。
|
||||||
- **📝 执行方案**:
|
- **📝 执行方案**:
|
||||||
1. 创建一个新脚本，例如 `tools/generate_synthetic_layouts.py`。
|
1. 创建一个新脚本，例如 `tools/generate_synthetic_layouts.py`。
|
||||||
2. 利用 `gdstk` 库 编写函数，程序化地生成包含不同尺寸、密度和类型标准单元的 GDSII 文件。
|
2. 利用 `gdstk` 库 编写函数，程序化地生成包含不同尺寸、密度和类型标准单元的 GDSII 文件。
|
||||||
3. 结合 `tools/layout2png.py` 的逻辑，将生成的版图批量转换为 PNG 图像，用于扩充训练集。
|
3. 结合 `tools/layout2png.py` 的逻辑，将生成的版图批量转换为 PNG 图像，用于扩充训练集。
|
||||||
|
|
||||||
|
- [ ] **基于扩散生成的版图数据生成器（研究型）**
|
||||||
|
- **🎯 目标**: 使用扩散模型（Diffusion）生成具备“曼哈顿几何特性”的版图切片（raster PNG），作为现有程序化合成的补充来源，进一步提升数据多样性与风格覆盖。
|
||||||
|
- **📦 产物**:
|
||||||
|
- 推理脚本（计划）: `tools/diffusion/sample_layouts.py`
|
||||||
|
- 训练脚本（计划）: `tools/diffusion/train_layout_diffusion.py`
|
||||||
|
- 数据集打包与统计工具（计划）: `tools/diffusion/prepare_patch_dataset.py`
|
||||||
|
- **🧭 范围界定**:
|
||||||
|
- 优先生成单层的二值/灰度光栅图像（256–512 像素方形 patch）。
|
||||||
|
- 短期不追求多层/DRC 严格约束的工业可制造性；定位为数据增强来源，而非版图设计替代。
|
||||||
|
- **🛤️ 技术路线**:
|
||||||
|
- 路线 A（首选，工程落地快）: 基于 HuggingFace diffusers 的 Latent Diffusion/Stable Diffusion 微调；输入为 1 通道灰度（训练时复制到 3 通道或改 UNet 首层），输出为版图样式图像。
|
||||||
|
- 路线 B（结构引导）: 加入 ControlNet/T2I-Adapter 条件，如 Sobel/Canny/直方结构图、粗草图（Scribble）、程序化几何草图，以控制生成的总体连通性与直角占比。
|
||||||
|
- 路线 C（两阶段）: 先用程序化生成器输出“草图/骨架”（低细节），再用扩散模型进行“风格化/细化”。
|
||||||
|
- **🧱 数据表示与条件**:
|
||||||
|
- Raster 表示：PNG（二值/灰度），可预生成条件图：Sobel、Canny、距离变换、形态学骨架等。
|
||||||
|
- 条件输入建议：`[image (target-like), edge_map, skeleton]` 的任意子集；PoC 以 edge_map 为主。
|
||||||
|
- **🧪 训练配置（建议起点）**:
|
||||||
|
- 图像尺寸：256（PoC），后续 384/512。
|
||||||
|
- 批大小：8–16（依显存），学习率 1e-4，训练步数 100k–300k。
|
||||||
|
- 数据来源：`data/**/png` 聚合 + 程序合成数据 `data/synthetic/png`；采样时按风格/密度分层均衡。
|
||||||
|
- 预处理：随机裁剪非空 patch、二值阈值均衡、弱摄影增强（噪声/对比度）控制在小幅度范围。
|
||||||
|
- **🧰 推理与后处理**:
|
||||||
|
- 采样参数：采样步数 30–100、guidance scale 3–7、seed 固定以便复现。
|
||||||
|
- 后处理：Otsu/固定阈值二值化，形态学开闭/细化，断点连接（morphology bridge），可选矢量化（`gdstk` 轮廓化）回写 GDS。
|
||||||
|
- **📈 评估指标**:
|
||||||
|
- 结构统计对齐：水平/垂直边比例、连通组件面积分布、线宽分布、密度直方图与真实数据 KL 距离。
|
||||||
|
- 规则近似性：形态学开闭后碎片率、连通率、冗余孤立像素占比。
|
||||||
|
- 训练收益：将扩散样本混入 `train.py`，对 IoU/mAP/收敛轮数的提升幅度（与仅程序合成相比）。
|
||||||
|
- **🔌 与现有管线集成**:
|
||||||
|
- 在 `tools/synth_pipeline.py` 增加 `--use_diffusion` 或 `--diffusion_dir`，将扩散生成的 PNG 目录并入训练数据目录。
|
||||||
|
- 配置建议新增：
|
||||||
|
```yaml
|
||||||
|
synthetic:
|
||||||
|
diffusion:
|
||||||
|
enabled: false
|
||||||
|
png_dir: data/synthetic_diff/png
|
||||||
|
ratio: 0.1 # 与真实/程序合成的混采比例
|
||||||
|
```
|
||||||
|
- 预览与质检：重用 `tools/preview_dataset.py`，并用 `tools/validate_h_consistency.py` 跳过 H 检查（扩散输出无严格几何配对），改用结构统计工具（后续补充）。
|
||||||
|
- **🗓️ 里程碑**:
|
||||||
|
1. 第 1 周：数据准备与统计、PoC（预训练 SD + ControlNet-Edge 的小规模微调，256 尺寸）。
|
||||||
|
2. 第 2–3 周：扩大训练（≥50k patch），加入骨架/距离变换条件，完善后处理。
|
||||||
|
3. 第 4 周：与训练管线集成（混采/可视化），对比“仅程序合成 vs 程序合成+扩散”的增益。
|
||||||
|
4. 第 5 周：文档、示例权重与一键脚本（可选导出 ONNX/TensorRT 推理）。
|
||||||
|
- **⚠️ 风险与缓解**:
|
||||||
|
- 结构失真/非曼哈顿：增强条件约束（ControlNet），提高形态学后处理强度；两阶段（草图→细化）。
|
||||||
|
- 模式崩塌/多样性不足：分层采样、数据重采样、EMA、风格/密度条件编码。
|
||||||
|
- 训练数据不足：先用程序合成预训练，再混入少量真实数据微调。
|
||||||
|
- **📚 参考与依赖**:
|
||||||
|
- 依赖：`diffusers`, `transformers`, `accelerate`, `albumentations`, `opencv-python`, `gdstk`
|
||||||
|
- 参考：Latent Diffusion、Stable Diffusion、ControlNet、T2I-Adapter 等论文与开源实现
|
||||||
|
|
||||||
### 二、 模型架构 (Model Architecture)
|
### 二、 模型架构 (Model Architecture)
|
||||||
|
|
||||||
> *目标：提升模型的特征提取效率和精度，降低计算资源消耗。*
|
> *目标：提升模型的特征提取效率和精度，降低计算资源消耗。*
|
||||||
@@ -40,11 +128,19 @@
|
|||||||
- 代码：`models/rord.py`
|
- 代码：`models/rord.py`
|
||||||
- 基准：`tests/benchmark_backbones.py`
|
- 基准：`tests/benchmark_backbones.py`
|
||||||
- 文档：`docs/description/Backbone_FPN_Test_Change_Notes.md`, `docs/description/Performance_Benchmark.md`
|
- 文档：`docs/description/Backbone_FPN_Test_Change_Notes.md`, `docs/description/Performance_Benchmark.md`
|
||||||
- [ ] **集成注意力机制 (Attention Mechanism)**
|
- [x] **集成注意力机制 (Attention Mechanism)**
|
||||||
- **✔️ 价值**: 引导模型自动关注版图中的关键几何结构（如边角、交点），忽略大面积的空白或重复区域，提升特征质量。
|
- **✔️ 价值**: 引导模型关注关键几何结构、弱化冗余区域，提升特征质量与匹配稳定性。
|
||||||
- **📝 执行方案**:
|
- **✅ 当前进展（2025-10-20）**:
|
||||||
1. 寻找一个可靠的注意力模块实现，如 CBAM 或 SE-Net。
|
- 已集成可切换的注意力模块：`SE` 与 `CBAM`；支持通过 `model.attention.enabled/type/places` 配置开启与插入位置（`backbone_high`/`det_head`/`desc_head`）。
|
||||||
2. 在 `models/rord.py` 中，将该模块插入到 `self.backbone` 和两个 `head` 之间。
|
- 已完成 CPU A/B 基准（none/se/cbam，resnet34，places=backbone_high+desc_head），详见 `docs/description/Performance_Benchmark.md`；脚本：`tests/benchmark_attention.py`。
|
||||||
|
- **📝 后续动作**:
|
||||||
|
1. 扩展更多模块：ECA、SimAM、CoordAttention、SKNet，并保持统一接口与配置。
|
||||||
|
2. 进行插入位置消融（仅 backbone_high / det_head / desc_head / 组合），在 GPU 上复测速度与显存峰值。
|
||||||
|
3. 在真实数据上评估注意力开/关的 IoU/mAP 与收敛差异。
|
||||||
|
- **参考**:
|
||||||
|
- 代码：`models/rord.py`
|
||||||
|
- 基准：`tests/benchmark_attention.py`, `tests/benchmark_grid.py`
|
||||||
|
- 文档：`docs/description/Performance_Benchmark.md`
|
||||||
|
|
||||||
### 三、 训练与损失函数 (Training & Loss Function)
|
### 三、 训练与损失函数 (Training & Loss Function)
|
||||||
|
|
||||||
@@ -181,6 +277,24 @@
|
|||||||
--output-file export.md
|
--output-file export.md
|
||||||
```
|
```
|
||||||
|
|
||||||
|
### ✅ 三维基准对比（Backbone × Attention × Single/FPN）
|
||||||
|
|
||||||
|
- **文件**: `tests/benchmark_grid.py` ✅，JSON 输出：`benchmark_grid.json`
|
||||||
|
- **功能**:
|
||||||
|
- 遍历 `backbone × attention` 组合（当前：vgg16/resnet34/efficientnet_b0 × none/se/cbam）
|
||||||
|
- 统计单尺度与 FPN 前向的平均耗时与标准差
|
||||||
|
- 控制台摘要 + JSON 结果落盘
|
||||||
|
- **使用**:
|
||||||
|
```bash
|
||||||
|
PYTHONPATH=. uv run python tests/benchmark_grid.py \
|
||||||
|
--device cpu --image-size 512 --runs 3 \
|
||||||
|
--backbones vgg16 resnet34 efficientnet_b0 \
|
||||||
|
--attentions none se cbam \
|
||||||
|
--places backbone_high desc_head
|
||||||
|
```
|
||||||
|
- **结果**:
|
||||||
|
- 已将 CPU（512×512，runs=3）结果写入 `docs/description/Performance_Benchmark.md` 的“三维基准”表格，原始数据位于仓库根目录 `benchmark_grid.json`。
|
||||||
|
|
||||||
### 📚 新增文档
|
### 📚 新增文档
|
||||||
|
|
||||||
| 文档 | 大小 | 说明 |
|
| 文档 | 大小 | 说明 |
|
||||||
@@ -263,6 +377,8 @@
|
|||||||
| | 全面评估指标 | ✅ | 2025-10-19 |
|
| | 全面评估指标 | ✅ | 2025-10-19 |
|
||||||
| **新增工作** | 性能基准测试 | ✅ | 2025-10-20 |
|
| **新增工作** | 性能基准测试 | ✅ | 2025-10-20 |
|
||||||
| | TensorBoard 导出工具 | ✅ | 2025-10-20 |
|
| | TensorBoard 导出工具 | ✅ | 2025-10-20 |
|
||||||
|
| **二. 模型架构** | 注意力机制（SE/CBAM 基线） | ✅ | 2025-10-20 |
|
||||||
|
| **新增工作** | 三维基准对比（Backbone×Attention×Single/FPN） | ✅ | 2025-10-20 |
|
||||||
|
|
||||||
### 未完成的工作项（可选优化）
|
### 未完成的工作项（可选优化）
|
||||||
|
|
||||||
@@ -270,8 +386,52 @@
|
|||||||
|------|--------|--------|------|
|
|------|--------|--------|------|
|
||||||
| **一. 数据策略与增强** | 弹性变形增强 | 🟡 低 | 便利性增强 |
|
| **一. 数据策略与增强** | 弹性变形增强 | 🟡 低 | 便利性增强 |
|
||||||
| | 合成版图生成器 | 🟡 低 | 数据增强 |
|
| | 合成版图生成器 | 🟡 低 | 数据增强 |
|
||||||
| **二. 模型架构** | 现代骨干网络 | 🟠 中 | 性能优化 |
|
| | 基于扩散的版图生成器 | 🟠 中 | 研究型：引入结构条件与形态学后处理，作为数据多样性来源 |
|
||||||
| | 注意力机制 | 🟠 中 | 性能优化 |
|
|
||||||
|
---
|
||||||
|
|
||||||
|
## 扩散生成集成的实现说明（新增）
|
||||||
|
|
||||||
|
- 配置新增节点（已添加到 `configs/base_config.yaml`）:
|
||||||
|
```yaml
|
||||||
|
synthetic:
|
||||||
|
enabled: false
|
||||||
|
png_dir: data/synthetic/png
|
||||||
|
ratio: 0.0
|
||||||
|
diffusion:
|
||||||
|
enabled: false
|
||||||
|
png_dir: data/synthetic_diff/png
|
||||||
|
ratio: 0.0
|
||||||
|
```
|
||||||
|
|
||||||
|
- 训练混采（已实现于 `train.py`）:
|
||||||
|
- 支持三源混采：真实数据 + 程序合成 (`synthetic`) + 扩散合成 (`synthetic.diffusion`)。
|
||||||
|
- 目标比例：`real = 1 - (syn_ratio + diff_ratio)`；使用 `WeightedRandomSampler` 近似。
|
||||||
|
- 验证集仅使用真实数据，避免评估偏移。
|
||||||
|
|
||||||
|
- 一键管线扩展（已实现于 `tools/synth_pipeline.py`）:
|
||||||
|
- 新增 `--diffusion_dir` 参数：将指定目录的 PNG 并入配置文件的 `synthetic.diffusion.png_dir` 并开启 `enabled=true`。
|
||||||
|
- 不自动采样扩散图片（避免引入新依赖），仅做目录集成；后续可在该脚本中串联 `tools/diffusion/sample_layouts.py`。
|
||||||
|
|
||||||
|
- 新增脚本骨架（`tools/diffusion/`）:
|
||||||
|
- `prepare_patch_dataset.py`: 从现有 PNG 构建 patch 数据集与条件图（CLI 骨架 + TODO）。
|
||||||
|
- `train_layout_diffusion.py`: 微调扩散模型的训练脚本（CLI 骨架 + TODO）。
|
||||||
|
- `sample_layouts.py`: 使用已训练权重进行采样输出 PNG（CLI 骨架 + TODO）。
|
||||||
|
|
||||||
|
- 使用建议：
|
||||||
|
1) 将扩散采样得到的 PNG 放入某目录，例如 `data/synthetic_diff/png`。
|
||||||
|
2) 运行：
|
||||||
|
```bash
|
||||||
|
uv run python tools/synth_pipeline.py \
|
||||||
|
--out_root data/synthetic \
|
||||||
|
--num 200 --dpi 600 \
|
||||||
|
--config configs/base_config.yaml \
|
||||||
|
--ratio 0.3 \
|
||||||
|
--diffusion_dir data/synthetic_diff/png
|
||||||
|
```
|
||||||
|
3) 在 YAML 中按需设置 `synthetic.diffusion.ratio`（例如 0.1），训练时即自动按比例混采。
|
||||||
|
|
||||||
|
| **二. 模型架构** | 更多注意力模块（ECA/SimAM/CoordAttention/SKNet） | 🟠 中 | 扩展与消融 |
|
||||||
| **三. 训练与损失** | 损失加权自适应 | 🟠 中 | 训练优化 |
|
| **三. 训练与损失** | 损失加权自适应 | 🟠 中 | 训练优化 |
|
||||||
| | 困难样本采样 | 🟡 低 | 训练优化 |
|
| | 困难样本采样 | 🟡 低 | 训练优化 |
|
||||||
|
|
||||||
|
|||||||
@@ -82,6 +82,25 @@
|
|||||||
- [ ] 特征维度一致性检查
|
- [ ] 特征维度一致性检查
|
||||||
- [ ] GPU/CPU 切换测试
|
- [ ] GPU/CPU 切换测试
|
||||||
|
|
||||||
|
#### 2.3 基准与评估补充（来自 NextStep 2.1 未完项）
|
||||||
|
|
||||||
|
- [ ] GPU 环境 A/B 基准（速度/显存）
|
||||||
|
- [ ] 使用 `tests/benchmark_backbones.py` 在 GPU 上复现（20 次，512×512），记录 ms 与 VRAM
|
||||||
|
- [ ] 追加结果到 `docs/description/Performance_Benchmark.md`
|
||||||
|
|
||||||
|
- [ ] GPU 环境 Attention A/B 基准（速度/显存）
|
||||||
|
- [ ] 使用 `tests/benchmark_attention.py` 在 GPU 上复现（10 次，512×512），覆盖 `places` 组合（`backbone_high`/`det_head`/`desc_head`）
|
||||||
|
- [ ] 记录平均耗时与 VRAM 峰值，追加摘要到 `docs/description/Performance_Benchmark.md`
|
||||||
|
|
||||||
|
- [ ] 三维网格基准（Backbone × Attention × Single/FPN）
|
||||||
|
- [ ] 使用 `tests/benchmark_grid.py` 在 GPU 上跑最小矩阵（例如 3×3，runs=5）
|
||||||
|
- [ ] 将 JSON 存入 `results/benchmark_grid_YYYYMMDD.json`，在性能文档中追加表格摘要并链接 JSON
|
||||||
|
|
||||||
|
- [ ] 真实数据集精度评估（IoU/mAP 与收敛曲线）
|
||||||
|
- [ ] 固定数据与超参，训练 5 个 epoch，记录 loss 曲线
|
||||||
|
- [ ] 在验证集上评估 IoU/mAP，并与 vgg16 基线对比
|
||||||
|
- [ ] 形成对照表与初步结论
|
||||||
|
|
||||||
**验收标准**:
|
**验收标准**:
|
||||||
- [ ] 所有测试用例通过
|
- [ ] 所有测试用例通过
|
||||||
- [ ] 推理结果符合预期维度和范围
|
- [ ] 推理结果符合预期维度和范围
|
||||||
@@ -143,6 +162,12 @@
|
|||||||
- [ ] 日志查看方法
|
- [ ] 日志查看方法
|
||||||
- [ ] GPU 内存不足处理
|
- [ ] GPU 内存不足处理
|
||||||
|
|
||||||
|
#### 3.4 预训练权重加载摘要（来自 NextStep 2.1 未完项）
|
||||||
|
|
||||||
|
- [x] 在 `models/rord.py` 加载 `pretrained=true` 时，打印未命中层摘要
|
||||||
|
- [x] 记录：加载成功/跳过的层名数量
|
||||||
|
- [x] 提供简要输出（missing/unexpected keys，参数量统计）；实现：`models/rord.py::_summarize_pretrained_load`
|
||||||
|
|
||||||
#### 3.2 编写配置参数文档
|
#### 3.2 编写配置参数文档
|
||||||
|
|
||||||
- [ ] 创建 `docs/CONFIG.md`
|
- [ ] 创建 `docs/CONFIG.md`
|
||||||
|
|||||||
@@ -218,6 +218,41 @@
|
|||||||
|
|
||||||
---
|
---
|
||||||
|
|
||||||
|
### 4. 注意力机制集成（来自 NextStep 2.2）
|
||||||
|
|
||||||
|
**目标**: 在骨干高层与头部前集成 CBAM / SE，并量化收益
|
||||||
|
|
||||||
|
#### 4.1 模块实现与插桩
|
||||||
|
- [ ] 实现 `CBAM` 与 `SEBlock`（或迁移可靠实现）
|
||||||
|
- [ ] 在 `models/rord.py` 通过配置插拔：`attention.enabled/type/places`
|
||||||
|
- [ ] 确保 forward 尺寸不变，默认关闭可回退
|
||||||
|
|
||||||
|
#### 4.2 训练与评估
|
||||||
|
- [ ] 选择入选骨干为基线，分别开启 `cbam` 与 `se`
|
||||||
|
- [ ] 记录训练损失、验证 IoU/mAP、推理时延/显存
|
||||||
|
- [ ] 可选：导出可视化注意力图
|
||||||
|
|
||||||
|
**验收标准**:
|
||||||
|
- [ ] 训练稳定，无数值异常
|
||||||
|
- [ ] 指标不低于无注意力基线；若提升则量化收益
|
||||||
|
- [ ] 配置可一键关闭以回退
|
||||||
|
|
||||||
|
#### 4.3 扩展模块与插入位置消融
|
||||||
|
- [ ] 扩展更多注意力模块：ECA、SimAM、CoordAttention、SKNet
|
||||||
|
- [ ] 在 `models/rord.py` 实现统一接口与注册表
|
||||||
|
- [ ] 在 `configs/base_config.yaml` 增加可选项说明
|
||||||
|
- [ ] 插入位置消融
|
||||||
|
- [ ] 仅 `backbone_high` / 仅 `det_head` / 仅 `desc_head` / 组合
|
||||||
|
- [ ] 使用 `tests/benchmark_attention.py` 统一基准，记录 Single/FPN 时延与 VRAM
|
||||||
|
- [ ] 在 `docs/description/Performance_Benchmark.md` 增加“注意力插入位置”小节
|
||||||
|
|
||||||
|
**验收标准**:
|
||||||
|
- [ ] 所有新增模块 forward 通过，尺寸/类型与现有路径一致
|
||||||
|
- [ ] 基准结果可复现并写入文档
|
||||||
|
- [ ] 给出速度-精度权衡建议
|
||||||
|
|
||||||
|
---
|
||||||
|
|
||||||
## 🔄 实施流程
|
## 🔄 实施流程
|
||||||
|
|
||||||
### 第 1 周: 实验管理集成
|
### 第 1 周: 实验管理集成
|
||||||
|
|||||||
@@ -91,7 +91,12 @@ class RoRD(nn.Module):
|
|||||||
# 默认各层通道（VGG 对齐）
|
# 默认各层通道（VGG 对齐）
|
||||||
c2_ch, c3_ch, c4_ch = 128, 256, 512
|
c2_ch, c3_ch, c4_ch = 128, 256, 512
|
||||||
if backbone_name == "resnet34":
|
if backbone_name == "resnet34":
|
||||||
res = models.resnet34(weights=models.ResNet34_Weights.DEFAULT if pretrained else None)
|
# 构建骨干并按需手动加载权重，便于打印加载摘要
|
||||||
|
if pretrained:
|
||||||
|
res = models.resnet34(weights=None)
|
||||||
|
self._summarize_pretrained_load(res, models.ResNet34_Weights.DEFAULT, "resnet34")
|
||||||
|
else:
|
||||||
|
res = models.resnet34(weights=None)
|
||||||
self.backbone = nn.Sequential(
|
self.backbone = nn.Sequential(
|
||||||
res.conv1, res.bn1, res.relu, res.maxpool,
|
res.conv1, res.bn1, res.relu, res.maxpool,
|
||||||
res.layer1, res.layer2, res.layer3, res.layer4,
|
res.layer1, res.layer2, res.layer3, res.layer4,
|
||||||
@@ -102,14 +107,23 @@ class RoRD(nn.Module):
|
|||||||
# 选择 layer2/layer3/layer4 作为 C2/C3/C4
|
# 选择 layer2/layer3/layer4 作为 C2/C3/C4
|
||||||
c2_ch, c3_ch, c4_ch = 128, 256, 512
|
c2_ch, c3_ch, c4_ch = 128, 256, 512
|
||||||
elif backbone_name == "efficientnet_b0":
|
elif backbone_name == "efficientnet_b0":
|
||||||
eff = models.efficientnet_b0(weights=models.EfficientNet_B0_Weights.DEFAULT if pretrained else None)
|
if pretrained:
|
||||||
|
eff = models.efficientnet_b0(weights=None)
|
||||||
|
self._summarize_pretrained_load(eff, models.EfficientNet_B0_Weights.DEFAULT, "efficientnet_b0")
|
||||||
|
else:
|
||||||
|
eff = models.efficientnet_b0(weights=None)
|
||||||
self.backbone = eff.features
|
self.backbone = eff.features
|
||||||
self._backbone_raw = eff
|
self._backbone_raw = eff
|
||||||
out_channels_backbone = 1280
|
out_channels_backbone = 1280
|
||||||
# 选择 features[2]/[3]/[6] 作为 C2/C3/C4（约 24/40/192）
|
# 选择 features[2]/[3]/[6] 作为 C2/C3/C4（约 24/40/192）
|
||||||
c2_ch, c3_ch, c4_ch = 24, 40, 192
|
c2_ch, c3_ch, c4_ch = 24, 40, 192
|
||||||
else:
|
else:
|
||||||
vgg16_features = models.vgg16(weights=models.VGG16_Weights.DEFAULT if pretrained else None).features
|
if pretrained:
|
||||||
|
vgg = models.vgg16(weights=None)
|
||||||
|
self._summarize_pretrained_load(vgg, models.VGG16_Weights.DEFAULT, "vgg16")
|
||||||
|
else:
|
||||||
|
vgg = models.vgg16(weights=None)
|
||||||
|
vgg16_features = vgg.features
|
||||||
# VGG16 特征各阶段索引（conv & relu 层序列）
|
# VGG16 特征各阶段索引（conv & relu 层序列）
|
||||||
# relu2_2 索引 8，relu3_3 索引 15，relu4_3 索引 22
|
# relu2_2 索引 8，relu3_3 索引 15，relu4_3 索引 22
|
||||||
self.features = vgg16_features
|
self.features = vgg16_features
|
||||||
@@ -264,3 +278,32 @@ class RoRD(nn.Module):
|
|||||||
return c2, c3, c4
|
return c2, c3, c4
|
||||||
|
|
||||||
raise RuntimeError(f"Unsupported backbone for FPN: {self.backbone_name}")
|
raise RuntimeError(f"Unsupported backbone for FPN: {self.backbone_name}")
|
||||||
|
|
||||||
|
# --- Utils ---
|
||||||
|
def _summarize_pretrained_load(self, torch_model: nn.Module, weights_enum, arch_name: str) -> None:
|
||||||
|
"""手动加载 torchvision 预训练权重并打印加载摘要。
|
||||||
|
- 使用 strict=False 以兼容可能的键差异，打印 missing/unexpected keys。
|
||||||
|
- 输出参数量统计，便于快速核对加载情况。
|
||||||
|
"""
|
||||||
|
try:
|
||||||
|
state_dict = weights_enum.get_state_dict(progress=False)
|
||||||
|
except Exception:
|
||||||
|
# 回退：若权重枚举不支持 get_state_dict，则跳过摘要（通常已在构造器中加载）
|
||||||
|
print(f"[Pretrained] {arch_name}: skip summary (weights enum lacks get_state_dict)")
|
||||||
|
return
|
||||||
|
incompatible = torch_model.load_state_dict(state_dict, strict=False)
|
||||||
|
total_params = sum(p.numel() for p in torch_model.parameters())
|
||||||
|
trainable_params = sum(p.numel() for p in torch_model.parameters() if p.requires_grad)
|
||||||
|
missing = list(getattr(incompatible, 'missing_keys', []))
|
||||||
|
unexpected = list(getattr(incompatible, 'unexpected_keys', []))
|
||||||
|
try:
|
||||||
|
matched = len(state_dict) - len(unexpected)
|
||||||
|
except Exception:
|
||||||
|
matched = 0
|
||||||
|
print(f"[Pretrained] {arch_name}: ImageNet weights loaded (strict=False)")
|
||||||
|
print(f" params: total={total_params/1e6:.2f}M, trainable={trainable_params/1e6:.2f}M")
|
||||||
|
print(f" keys: matched≈{matched} | missing={len(missing)} | unexpected={len(unexpected)}")
|
||||||
|
if missing and len(missing) <= 10:
|
||||||
|
print(f" missing: {missing}")
|
||||||
|
if unexpected and len(unexpected) <= 10:
|
||||||
|
print(f" unexpected: {unexpected}")
|
||||||
91
tests/benchmark_attention.py
Normal file
91
tests/benchmark_attention.py
Normal file
@@ -0,0 +1,91 @@
|
|||||||
|
"""
|
||||||
|
注意力模块 A/B 基准测试
|
||||||
|
|
||||||
|
目的：在相同骨干与输入下，对比注意力开/关（none/se/cbam）在单尺度与 FPN 前向的耗时差异；可选指定插入位置。
|
||||||
|
|
||||||
|
示例：
|
||||||
|
PYTHONPATH=. uv run python tests/benchmark_attention.py --device cpu --image-size 512 --runs 10 --backbone resnet34 --places backbone_high desc_head
|
||||||
|
"""
|
||||||
|
from __future__ import annotations
|
||||||
|
|
||||||
|
import argparse
|
||||||
|
import time
|
||||||
|
from typing import Dict, List
|
||||||
|
|
||||||
|
import numpy as np
|
||||||
|
import torch
|
||||||
|
|
||||||
|
from models.rord import RoRD
|
||||||
|
|
||||||
|
|
||||||
|
def bench_once(model: torch.nn.Module, x: torch.Tensor, fpn: bool = False) -> float:
|
||||||
|
if torch.cuda.is_available() and x.is_cuda:
|
||||||
|
torch.cuda.synchronize()
|
||||||
|
t0 = time.time()
|
||||||
|
with torch.inference_mode():
|
||||||
|
_ = model(x, return_pyramid=fpn)
|
||||||
|
if torch.cuda.is_available() and x.is_cuda:
|
||||||
|
torch.cuda.synchronize()
|
||||||
|
return (time.time() - t0) * 1000.0
|
||||||
|
|
||||||
|
|
||||||
|
def build_model(backbone: str, attention_type: str, places: List[str], device: torch.device) -> RoRD:
|
||||||
|
cfg = type("cfg", (), {
|
||||||
|
"model": type("m", (), {
|
||||||
|
"backbone": type("b", (), {"name": backbone, "pretrained": False})(),
|
||||||
|
"attention": type("a", (), {"enabled": attention_type != "none", "type": attention_type, "places": places})(),
|
||||||
|
})()
|
||||||
|
})()
|
||||||
|
model = RoRD(cfg=cfg).to(device)
|
||||||
|
model.eval()
|
||||||
|
return model
|
||||||
|
|
||||||
|
|
||||||
|
def run_suite(backbone: str, places: List[str], device: torch.device, image_size: int, runs: int) -> List[Dict[str, float]]:
|
||||||
|
x = torch.randn(1, 3, image_size, image_size, device=device)
|
||||||
|
results: List[Dict[str, float]] = []
|
||||||
|
for attn in ["none", "se", "cbam"]:
|
||||||
|
model = build_model(backbone, attn, places, device)
|
||||||
|
# warmup
|
||||||
|
for _ in range(3):
|
||||||
|
_ = model(x, return_pyramid=False)
|
||||||
|
_ = model(x, return_pyramid=True)
|
||||||
|
# single
|
||||||
|
t_list_single = [bench_once(model, x, fpn=False) for _ in range(runs)]
|
||||||
|
# fpn
|
||||||
|
t_list_fpn = [bench_once(model, x, fpn=True) for _ in range(runs)]
|
||||||
|
results.append({
|
||||||
|
"backbone": backbone,
|
||||||
|
"attention": attn,
|
||||||
|
"places": ",".join(places) if places else "-",
|
||||||
|
"single_ms_mean": float(np.mean(t_list_single)),
|
||||||
|
"single_ms_std": float(np.std(t_list_single)),
|
||||||
|
"fpn_ms_mean": float(np.mean(t_list_fpn)),
|
||||||
|
"fpn_ms_std": float(np.std(t_list_fpn)),
|
||||||
|
"runs": int(runs),
|
||||||
|
})
|
||||||
|
return results
|
||||||
|
|
||||||
|
|
||||||
|
def main():
|
||||||
|
parser = argparse.ArgumentParser(description="RoRD 注意力模块 A/B 基准")
|
||||||
|
parser.add_argument("--backbone", type=str, default="resnet34", choices=["vgg16","resnet34","efficientnet_b0"], help="骨干")
|
||||||
|
parser.add_argument("--places", nargs="*", default=["backbone_high"], help="插入位置：backbone_high det_head desc_head")
|
||||||
|
parser.add_argument("--image-size", type=int, default=512, help="输入尺寸")
|
||||||
|
parser.add_argument("--runs", type=int, default=10, help="重复次数")
|
||||||
|
parser.add_argument("--device", type=str, default="cpu", help="cuda 或 cpu")
|
||||||
|
args = parser.parse_args()
|
||||||
|
|
||||||
|
device = torch.device(args.device if torch.cuda.is_available() or args.device == "cpu" else "cpu")
|
||||||
|
results = run_suite(args.backbone, args.places, device, args.image_size, args.runs)
|
||||||
|
|
||||||
|
# 简要打印
|
||||||
|
print("\n===== Attention A/B Summary =====")
|
||||||
|
for r in results:
|
||||||
|
print(f"{r['backbone']:<14} attn={r['attention']:<5} places={r['places']:<24} "
|
||||||
|
f"single {r['single_ms_mean']:.2f}±{r['single_ms_std']:.2f} | "
|
||||||
|
f"fpn {r['fpn_ms_mean']:.2f}±{r['fpn_ms_std']:.2f} ms")
|
||||||
|
|
||||||
|
|
||||||
|
if __name__ == "__main__":
|
||||||
|
main()
|
||||||
96
tests/benchmark_grid.py
Normal file
96
tests/benchmark_grid.py
Normal file
@@ -0,0 +1,96 @@
|
|||||||
|
"""
|
||||||
|
三维基准对比：Backbone × Attention × (SingleMean / FPNMean)
|
||||||
|
|
||||||
|
示例：
|
||||||
|
PYTHONPATH=. uv run python tests/benchmark_grid.py --device cpu --image-size 512 --runs 5 \
|
||||||
|
--backbones vgg16 resnet34 efficientnet_b0 --attentions none se cbam --places backbone_high
|
||||||
|
"""
|
||||||
|
from __future__ import annotations
|
||||||
|
|
||||||
|
import argparse
|
||||||
|
import json
|
||||||
|
import time
|
||||||
|
from typing import Dict, List
|
||||||
|
|
||||||
|
import numpy as np
|
||||||
|
import torch
|
||||||
|
|
||||||
|
from models.rord import RoRD
|
||||||
|
|
||||||
|
|
||||||
|
def bench_once(model: torch.nn.Module, x: torch.Tensor, fpn: bool = False) -> float:
|
||||||
|
if torch.cuda.is_available() and x.is_cuda:
|
||||||
|
torch.cuda.synchronize()
|
||||||
|
t0 = time.time()
|
||||||
|
with torch.inference_mode():
|
||||||
|
_ = model(x, return_pyramid=fpn)
|
||||||
|
if torch.cuda.is_available() and x.is_cuda:
|
||||||
|
torch.cuda.synchronize()
|
||||||
|
return (time.time() - t0) * 1000.0
|
||||||
|
|
||||||
|
|
||||||
|
def build_model(backbone: str, attention: str, places: List[str], device: torch.device) -> RoRD:
|
||||||
|
cfg = type("cfg", (), {
|
||||||
|
"model": type("m", (), {
|
||||||
|
"backbone": type("b", (), {"name": backbone, "pretrained": False})(),
|
||||||
|
"attention": type("a", (), {"enabled": attention != "none", "type": attention, "places": places})(),
|
||||||
|
})()
|
||||||
|
})()
|
||||||
|
model = RoRD(cfg=cfg).to(device)
|
||||||
|
model.eval()
|
||||||
|
return model
|
||||||
|
|
||||||
|
|
||||||
|
def run_grid(backbones: List[str], attentions: List[str], places: List[str], device: torch.device, image_size: int, runs: int) -> List[Dict[str, float]]:
|
||||||
|
x = torch.randn(1, 3, image_size, image_size, device=device)
|
||||||
|
rows: List[Dict[str, float]] = []
|
||||||
|
for bk in backbones:
|
||||||
|
for attn in attentions:
|
||||||
|
model = build_model(bk, attn, places, device)
|
||||||
|
# warmup
|
||||||
|
for _ in range(3):
|
||||||
|
_ = model(x, return_pyramid=False)
|
||||||
|
_ = model(x, return_pyramid=True)
|
||||||
|
# bench
|
||||||
|
t_single = [bench_once(model, x, fpn=False) for _ in range(runs)]
|
||||||
|
t_fpn = [bench_once(model, x, fpn=True) for _ in range(runs)]
|
||||||
|
rows.append({
|
||||||
|
"backbone": bk,
|
||||||
|
"attention": attn,
|
||||||
|
"places": ",".join(places) if places else "-",
|
||||||
|
"single_ms_mean": float(np.mean(t_single)),
|
||||||
|
"single_ms_std": float(np.std(t_single)),
|
||||||
|
"fpn_ms_mean": float(np.mean(t_fpn)),
|
||||||
|
"fpn_ms_std": float(np.std(t_fpn)),
|
||||||
|
"runs": int(runs),
|
||||||
|
})
|
||||||
|
return rows
|
||||||
|
|
||||||
|
|
||||||
|
def main():
|
||||||
|
parser = argparse.ArgumentParser(description="三维基准：Backbone × Attention × (Single/FPN)")
|
||||||
|
parser.add_argument("--backbones", nargs="*", default=["vgg16","resnet34","efficientnet_b0"], help="骨干列表")
|
||||||
|
parser.add_argument("--attentions", nargs="*", default=["none","se","cbam"], help="注意力列表")
|
||||||
|
parser.add_argument("--places", nargs="*", default=["backbone_high"], help="插入位置")
|
||||||
|
parser.add_argument("--image-size", type=int, default=512)
|
||||||
|
parser.add_argument("--runs", type=int, default=5)
|
||||||
|
parser.add_argument("--device", type=str, default="cpu")
|
||||||
|
parser.add_argument("--json-out", type=str, default="benchmark_grid.json")
|
||||||
|
args = parser.parse_args()
|
||||||
|
|
||||||
|
device = torch.device(args.device if torch.cuda.is_available() or args.device == "cpu" else "cpu")
|
||||||
|
rows = run_grid(args.backbones, args.attentions, args.places, device, args.image_size, args.runs)
|
||||||
|
|
||||||
|
# 打印简表
|
||||||
|
print("\n===== Grid Summary (Backbone × Attention) =====")
|
||||||
|
for r in rows:
|
||||||
|
print(f"{r['backbone']:<14} attn={r['attention']:<5} places={r['places']:<16} single {r['single_ms_mean']:.2f} | fpn {r['fpn_ms_mean']:.2f} ms")
|
||||||
|
|
||||||
|
# 保存 JSON
|
||||||
|
with open(args.json_out, 'w') as f:
|
||||||
|
json.dump(rows, f, indent=2)
|
||||||
|
print(f"Saved: {args.json_out}")
|
||||||
|
|
||||||
|
|
||||||
|
if __name__ == "__main__":
|
||||||
|
main()
|
||||||
46
tools/diffusion/prepare_patch_dataset.py
Normal file
46
tools/diffusion/prepare_patch_dataset.py
Normal file
@@ -0,0 +1,46 @@
|
|||||||
|
#!/usr/bin/env python3
|
||||||
|
"""
|
||||||
|
Prepare raster patch dataset and optional condition maps for diffusion training.
|
||||||
|
|
||||||
|
Planned inputs:
|
||||||
|
- --src_dirs: one or more directories containing PNG layout images
|
||||||
|
- --out_dir: output root for images/ and conditions/
|
||||||
|
- --size: patch size (e.g., 256)
|
||||||
|
- --stride: sliding stride for patch extraction
|
||||||
|
- --min_fg_ratio: minimum foreground ratio to keep a patch (0-1)
|
||||||
|
- --make_conditions: flags to generate edge/skeleton/distance maps
|
||||||
|
|
||||||
|
Current status: CLI skeleton and TODOs only.
|
||||||
|
"""
|
||||||
|
from __future__ import annotations
|
||||||
|
|
||||||
|
import argparse
|
||||||
|
from pathlib import Path
|
||||||
|
|
||||||
|
|
||||||
|
def main() -> None:
|
||||||
|
parser = argparse.ArgumentParser(description="Prepare patch dataset for diffusion training (skeleton)")
|
||||||
|
parser.add_argument("--src_dirs", type=str, nargs="+", help="Source PNG dirs for layouts")
|
||||||
|
parser.add_argument("--out_dir", type=str, required=True, help="Output root directory")
|
||||||
|
parser.add_argument("--size", type=int, default=256, help="Patch size")
|
||||||
|
parser.add_argument("--stride", type=int, default=256, help="Patch stride")
|
||||||
|
parser.add_argument("--min_fg_ratio", type=float, default=0.02, help="Min foreground ratio to keep a patch")
|
||||||
|
parser.add_argument("--make_edge", action="store_true", help="Generate edge map conditions (e.g., Sobel/Canny)")
|
||||||
|
parser.add_argument("--make_skeleton", action="store_true", help="Generate morphological skeleton condition")
|
||||||
|
parser.add_argument("--make_dist", action="store_true", help="Generate distance transform condition")
|
||||||
|
args = parser.parse_args()
|
||||||
|
|
||||||
|
out_root = Path(args.out_dir)
|
||||||
|
out_root.mkdir(parents=True, exist_ok=True)
|
||||||
|
(out_root / "images").mkdir(exist_ok=True)
|
||||||
|
(out_root / "conditions").mkdir(exist_ok=True)
|
||||||
|
|
||||||
|
# TODO: implement extraction loop over src_dirs, crop patches, filter by min_fg_ratio,
|
||||||
|
# and save into images/; generate optional condition maps into conditions/ mirroring filenames.
|
||||||
|
# Keep file naming consistent: images/xxx.png, conditions/xxx_edge.png, etc.
|
||||||
|
|
||||||
|
print("[TODO] Implement patch extraction and condition map generation.")
|
||||||
|
|
||||||
|
|
||||||
|
if __name__ == "__main__":
|
||||||
|
main()
|
||||||
38
tools/diffusion/sample_layouts.py
Normal file
38
tools/diffusion/sample_layouts.py
Normal file
@@ -0,0 +1,38 @@
|
|||||||
|
#!/usr/bin/env python3
|
||||||
|
"""
|
||||||
|
Sample layout patches using a trained diffusion model (skeleton).
|
||||||
|
|
||||||
|
Outputs raster PNGs into a target directory compatible with current training pipeline (no H pairing).
|
||||||
|
|
||||||
|
Current status: CLI skeleton and TODOs only.
|
||||||
|
"""
|
||||||
|
from __future__ import annotations
|
||||||
|
|
||||||
|
import argparse
|
||||||
|
from pathlib import Path
|
||||||
|
|
||||||
|
|
||||||
|
def main() -> None:
|
||||||
|
parser = argparse.ArgumentParser(description="Sample layout patches from diffusion model (skeleton)")
|
||||||
|
parser.add_argument("--ckpt", type=str, required=True, help="Path to trained diffusion checkpoint or HF repo id")
|
||||||
|
parser.add_argument("--out_dir", type=str, required=True, help="Directory to write sampled PNGs")
|
||||||
|
parser.add_argument("--num", type=int, default=200)
|
||||||
|
parser.add_argument("--image_size", type=int, default=256)
|
||||||
|
parser.add_argument("--guidance", type=float, default=5.0)
|
||||||
|
parser.add_argument("--steps", type=int, default=50)
|
||||||
|
parser.add_argument("--seed", type=int, default=42)
|
||||||
|
parser.add_argument("--cond_dir", type=str, default=None, help="Optional condition maps directory")
|
||||||
|
parser.add_argument("--cond_types", type=str, nargs="*", default=None, help="e.g., edge skeleton dist")
|
||||||
|
args = parser.parse_args()
|
||||||
|
|
||||||
|
out_dir = Path(args.out_dir)
|
||||||
|
out_dir.mkdir(parents=True, exist_ok=True)
|
||||||
|
|
||||||
|
# TODO: load pipeline from ckpt, set scheduler, handle conditions if provided,
|
||||||
|
# sample args.num images, save as PNG files into out_dir.
|
||||||
|
|
||||||
|
print("[TODO] Implement diffusion sampling and PNG saving.")
|
||||||
|
|
||||||
|
|
||||||
|
if __name__ == "__main__":
|
||||||
|
main()
|
||||||
37
tools/diffusion/train_layout_diffusion.py
Normal file
37
tools/diffusion/train_layout_diffusion.py
Normal file
@@ -0,0 +1,37 @@
|
|||||||
|
#!/usr/bin/env python3
|
||||||
|
"""
|
||||||
|
Train a diffusion model for layout patch generation (skeleton).
|
||||||
|
|
||||||
|
Planned: fine-tune Stable Diffusion (or Latent Diffusion) with optional ControlNet edge/skeleton conditions.
|
||||||
|
|
||||||
|
Dependencies to consider: diffusers, transformers, accelerate, torch, torchvision, opencv-python.
|
||||||
|
|
||||||
|
Current status: CLI skeleton and TODOs only.
|
||||||
|
"""
|
||||||
|
from __future__ import annotations
|
||||||
|
|
||||||
|
import argparse
|
||||||
|
|
||||||
|
|
||||||
|
def main() -> None:
|
||||||
|
parser = argparse.ArgumentParser(description="Train diffusion model for layout patches (skeleton)")
|
||||||
|
parser.add_argument("--data_dir", type=str, required=True, help="Prepared dataset root (images/ + conditions/)")
|
||||||
|
parser.add_argument("--output_dir", type=str, required=True, help="Checkpoint output directory")
|
||||||
|
parser.add_argument("--image_size", type=int, default=256)
|
||||||
|
parser.add_argument("--batch_size", type=int, default=8)
|
||||||
|
parser.add_argument("--lr", type=float, default=1e-4)
|
||||||
|
parser.add_argument("--max_steps", type=int, default=100000)
|
||||||
|
parser.add_argument("--use_controlnet", action="store_true", help="Train with ControlNet conditioning")
|
||||||
|
parser.add_argument("--condition_types", type=str, nargs="*", default=["edge"], help="e.g., edge skeleton dist")
|
||||||
|
args = parser.parse_args()
|
||||||
|
|
||||||
|
# TODO: implement dataset/dataloader (images and optional conditions)
|
||||||
|
# TODO: load base pipeline (Stable Diffusion or Latent Diffusion) and optionally ControlNet
|
||||||
|
# TODO: set up optimizer, LR schedule, EMA, gradient accumulation, and run training loop
|
||||||
|
# TODO: save periodic checkpoints to output_dir
|
||||||
|
|
||||||
|
print("[TODO] Implement diffusion training loop and checkpoints.")
|
||||||
|
|
||||||
|
|
||||||
|
if __name__ == "__main__":
|
||||||
|
main()
|
||||||
90
tools/generate_synthetic_layouts.py
Normal file
90
tools/generate_synthetic_layouts.py
Normal file
@@ -0,0 +1,90 @@
|
|||||||
|
#!/usr/bin/env python3
|
||||||
|
"""
|
||||||
|
Programmatic synthetic IC layout generator using gdstk.
|
||||||
|
Generates GDS files with simple standard-cell-like patterns, wires, and vias.
|
||||||
|
"""
|
||||||
|
from __future__ import annotations
|
||||||
|
|
||||||
|
import argparse
|
||||||
|
from pathlib import Path
|
||||||
|
import random
|
||||||
|
|
||||||
|
import gdstk
|
||||||
|
|
||||||
|
|
||||||
|
def build_standard_cell(cell_name: str, rng: random.Random, layer: int = 1, datatype: int = 0) -> gdstk.Cell:
|
||||||
|
cell = gdstk.Cell(cell_name)
|
||||||
|
# Basic cell body
|
||||||
|
w = rng.uniform(0.8, 2.0)
|
||||||
|
h = rng.uniform(1.6, 4.0)
|
||||||
|
rect = gdstk.rectangle((0, 0), (w, h), layer=layer, datatype=datatype)
|
||||||
|
cell.add(rect)
|
||||||
|
# Poly fingers
|
||||||
|
nf = rng.randint(1, 4)
|
||||||
|
pitch = w / (nf + 1)
|
||||||
|
for i in range(1, nf + 1):
|
||||||
|
x = i * pitch
|
||||||
|
poly = gdstk.rectangle((x - 0.05, 0), (x + 0.05, h), layer=layer + 1, datatype=datatype)
|
||||||
|
cell.add(poly)
|
||||||
|
# Contact/vias
|
||||||
|
for i in range(rng.randint(2, 6)):
|
||||||
|
vx = rng.uniform(0.1, w - 0.1)
|
||||||
|
vy = rng.uniform(0.1, h - 0.1)
|
||||||
|
via = gdstk.rectangle((vx - 0.05, vy - 0.05), (vx + 0.05, vy + 0.05), layer=layer + 2, datatype=datatype)
|
||||||
|
cell.add(via)
|
||||||
|
return cell
|
||||||
|
|
||||||
|
|
||||||
|
def generate_layout(out_path: Path, width: float, height: float, seed: int, rows: int, cols: int, density: float):
|
||||||
|
rng = random.Random(seed)
|
||||||
|
lib = gdstk.Library()
|
||||||
|
top = gdstk.Cell("TOP")
|
||||||
|
|
||||||
|
# Create a few standard cell variants
|
||||||
|
variants = [build_standard_cell(f"SC_{i}", rng, layer=1) for i in range(4)]
|
||||||
|
|
||||||
|
# Place instances in a grid with random skips based on density
|
||||||
|
x_pitch = width / cols
|
||||||
|
y_pitch = height / rows
|
||||||
|
for r in range(rows):
|
||||||
|
for c in range(cols):
|
||||||
|
if rng.random() > density:
|
||||||
|
continue
|
||||||
|
cell = rng.choice(variants)
|
||||||
|
dx = c * x_pitch + rng.uniform(0.0, 0.1 * x_pitch)
|
||||||
|
dy = r * y_pitch + rng.uniform(0.0, 0.1 * y_pitch)
|
||||||
|
ref = gdstk.Reference(cell, (dx, dy))
|
||||||
|
top.add(ref)
|
||||||
|
|
||||||
|
lib.add(*variants)
|
||||||
|
lib.add(top)
|
||||||
|
lib.write_gds(str(out_path))
|
||||||
|
|
||||||
|
|
||||||
|
def main():
|
||||||
|
parser = argparse.ArgumentParser(description="Generate synthetic IC layouts (GDS)")
|
||||||
|
parser.add_argument("--out-dir", type=str, default="data/synthetic/gds")
|
||||||
|
parser.add_argument("--out_dir", dest="out_dir", type=str, help="Alias of --out-dir")
|
||||||
|
parser.add_argument("--num-samples", type=int, default=10)
|
||||||
|
parser.add_argument("--num", dest="num_samples", type=int, help="Alias of --num-samples")
|
||||||
|
parser.add_argument("--seed", type=int, default=42)
|
||||||
|
parser.add_argument("--width", type=float, default=200.0)
|
||||||
|
parser.add_argument("--height", type=float, default=200.0)
|
||||||
|
parser.add_argument("--rows", type=int, default=10)
|
||||||
|
parser.add_argument("--cols", type=int, default=10)
|
||||||
|
parser.add_argument("--density", type=float, default=0.5)
|
||||||
|
|
||||||
|
args = parser.parse_args()
|
||||||
|
out_dir = Path(args.out_dir)
|
||||||
|
out_dir.mkdir(parents=True, exist_ok=True)
|
||||||
|
|
||||||
|
rng = random.Random(args.seed)
|
||||||
|
for i in range(args.num_samples):
|
||||||
|
sample_seed = rng.randint(0, 2**31 - 1)
|
||||||
|
out_path = out_dir / f"chip_{i:06d}.gds"
|
||||||
|
generate_layout(out_path, args.width, args.height, sample_seed, args.rows, args.cols, args.density)
|
||||||
|
print(f"[OK] Generated {out_path}")
|
||||||
|
|
||||||
|
|
||||||
|
if __name__ == "__main__":
|
||||||
|
main()
|
||||||
160
tools/layout2png.py
Normal file
160
tools/layout2png.py
Normal file
@@ -0,0 +1,160 @@
|
|||||||
|
#!/usr/bin/env python3
|
||||||
|
"""
|
||||||
|
Batch convert GDS to PNG.
|
||||||
|
|
||||||
|
Priority:
|
||||||
|
1) Use KLayout in headless batch mode (most accurate view fidelity for IC layouts).
|
||||||
|
2) Fallback to gdstk(read) -> write SVG -> cairosvg to PNG (no KLayout dependency at runtime).
|
||||||
|
"""
|
||||||
|
from __future__ import annotations
|
||||||
|
|
||||||
|
import argparse
|
||||||
|
from pathlib import Path
|
||||||
|
import subprocess
|
||||||
|
import sys
|
||||||
|
import tempfile
|
||||||
|
|
||||||
|
import cairosvg
|
||||||
|
|
||||||
|
|
||||||
|
def klayout_convert(gds_path: Path, png_path: Path, dpi: int, layermap: str | None = None, line_width: int | None = None, bgcolor: str | None = None) -> bool:
|
||||||
|
"""Render using KLayout by invoking a temporary Python macro with paths embedded."""
|
||||||
|
# Prepare optional display config code
|
||||||
|
layer_cfg_code = ""
|
||||||
|
if layermap:
|
||||||
|
# layermap format: "LAYER/DATATYPE:#RRGGBB,..."
|
||||||
|
layer_cfg_code += "lprops = pya.LayerPropertiesNode()\n"
|
||||||
|
for spec in layermap.split(","):
|
||||||
|
spec = spec.strip()
|
||||||
|
if not spec:
|
||||||
|
continue
|
||||||
|
try:
|
||||||
|
ld, color = spec.split(":")
|
||||||
|
layer_s, datatype_s = ld.split("/")
|
||||||
|
color = color.strip()
|
||||||
|
layer_cfg_code += (
|
||||||
|
"lp = pya.LayerPropertiesNode()\n"
|
||||||
|
f"lp.layer = int({int(layer_s)})\n"
|
||||||
|
f"lp.datatype = int({int(datatype_s)})\n"
|
||||||
|
f"lp.fill_color = pya.Color.from_string('{color}')\n"
|
||||||
|
f"lp.frame_color = pya.Color.from_string('{color}')\n"
|
||||||
|
"lprops.insert(lp)\n"
|
||||||
|
)
|
||||||
|
except Exception:
|
||||||
|
# Ignore malformed entries
|
||||||
|
continue
|
||||||
|
layer_cfg_code += "cv.set_layer_properties(lprops)\n"
|
||||||
|
|
||||||
|
line_width_code = ""
|
||||||
|
if line_width is not None:
|
||||||
|
line_width_code = f"cv.set_config('default-draw-line-width', '{int(line_width)}')\n"
|
||||||
|
|
||||||
|
bg_code = ""
|
||||||
|
if bgcolor:
|
||||||
|
bg_code = f"cv.set_config('background-color', '{bgcolor}')\n"
|
||||||
|
|
||||||
|
script = f"""
|
||||||
|
import pya
|
||||||
|
ly = pya.Layout()
|
||||||
|
ly.read(r"{gds_path}")
|
||||||
|
cv = pya.LayoutView()
|
||||||
|
cv.load_layout(ly, 0)
|
||||||
|
cv.max_hier_levels = 20
|
||||||
|
{bg_code}
|
||||||
|
{line_width_code}
|
||||||
|
{layer_cfg_code}
|
||||||
|
cv.zoom_fit()
|
||||||
|
cv.save_image(r"{png_path}", {dpi}, 0)
|
||||||
|
"""
|
||||||
|
try:
|
||||||
|
with tempfile.NamedTemporaryFile(mode="w", suffix=".py", delete=False) as tf:
|
||||||
|
tf.write(script)
|
||||||
|
tf.flush()
|
||||||
|
macro_path = Path(tf.name)
|
||||||
|
# Run klayout in batch mode
|
||||||
|
res = subprocess.run(["klayout", "-zz", "-b", "-r", str(macro_path)], check=False, capture_output=True, text=True)
|
||||||
|
ok = res.returncode == 0 and png_path.exists()
|
||||||
|
if not ok:
|
||||||
|
# Print stderr for visibility when running manually
|
||||||
|
if res.stderr:
|
||||||
|
sys.stderr.write(res.stderr)
|
||||||
|
try:
|
||||||
|
macro_path.unlink(missing_ok=True) # type: ignore[arg-type]
|
||||||
|
except Exception:
|
||||||
|
pass
|
||||||
|
return ok
|
||||||
|
except FileNotFoundError:
|
||||||
|
# klayout command not found
|
||||||
|
return False
|
||||||
|
except Exception:
|
||||||
|
return False
|
||||||
|
|
||||||
|
|
||||||
|
def gdstk_fallback(gds_path: Path, png_path: Path, dpi: int) -> bool:
|
||||||
|
"""Fallback path: use gdstk to read GDS and write SVG, then cairosvg to PNG.
|
||||||
|
Note: This may differ visually from KLayout depending on layers/styles.
|
||||||
|
"""
|
||||||
|
try:
|
||||||
|
import gdstk # local import to avoid import cost when not needed
|
||||||
|
svg_path = png_path.with_suffix(".svg")
|
||||||
|
lib = gdstk.read_gds(str(gds_path))
|
||||||
|
tops = lib.top_level()
|
||||||
|
if not tops:
|
||||||
|
return False
|
||||||
|
# Combine tops into a single temporary cell for rendering
|
||||||
|
cell = tops[0]
|
||||||
|
# gdstk Cell has write_svg in recent versions
|
||||||
|
try:
|
||||||
|
cell.write_svg(str(svg_path)) # type: ignore[attr-defined]
|
||||||
|
except Exception:
|
||||||
|
# Older gdstk: write_svg available on Library
|
||||||
|
try:
|
||||||
|
lib.write_svg(str(svg_path)) # type: ignore[attr-defined]
|
||||||
|
except Exception:
|
||||||
|
return False
|
||||||
|
# Convert SVG to PNG
|
||||||
|
cairosvg.svg2png(url=str(svg_path), write_to=str(png_path), dpi=dpi)
|
||||||
|
try:
|
||||||
|
svg_path.unlink()
|
||||||
|
except Exception:
|
||||||
|
pass
|
||||||
|
return True
|
||||||
|
except Exception:
|
||||||
|
return False
|
||||||
|
|
||||||
|
|
||||||
|
def main():
|
||||||
|
parser = argparse.ArgumentParser(description="Convert GDS files to PNG")
|
||||||
|
parser.add_argument("--in", dest="in_dir", type=str, required=True, help="Input directory containing .gds files")
|
||||||
|
parser.add_argument("--out", dest="out_dir", type=str, required=True, help="Output directory to place .png files")
|
||||||
|
parser.add_argument("--dpi", type=int, default=600, help="Output resolution in DPI for rasterization")
|
||||||
|
parser.add_argument("--layermap", type=str, default=None, help="Layer color map, e.g. '1/0:#00FF00,2/0:#FF0000'")
|
||||||
|
parser.add_argument("--line_width", type=int, default=None, help="Default draw line width in pixels for KLayout display")
|
||||||
|
parser.add_argument("--bgcolor", type=str, default=None, help="Background color, e.g. '#000000' or 'black'")
|
||||||
|
|
||||||
|
args = parser.parse_args()
|
||||||
|
in_dir = Path(args.in_dir)
|
||||||
|
out_dir = Path(args.out_dir)
|
||||||
|
out_dir.mkdir(parents=True, exist_ok=True)
|
||||||
|
|
||||||
|
gds_files = sorted(in_dir.glob("*.gds"))
|
||||||
|
if not gds_files:
|
||||||
|
print(f"[WARN] No GDS files found in {in_dir}")
|
||||||
|
return
|
||||||
|
|
||||||
|
ok_cnt = 0
|
||||||
|
for gds in gds_files:
|
||||||
|
png_path = out_dir / (gds.stem + ".png")
|
||||||
|
ok = klayout_convert(gds, png_path, args.dpi, layermap=args.layermap, line_width=args.line_width, bgcolor=args.bgcolor)
|
||||||
|
if not ok:
|
||||||
|
ok = gdstk_fallback(gds, png_path, args.dpi)
|
||||||
|
if ok:
|
||||||
|
ok_cnt += 1
|
||||||
|
print(f"[OK] {gds.name} -> {png_path}")
|
||||||
|
else:
|
||||||
|
print(f"[FAIL] {gds.name}")
|
||||||
|
print(f"Done. {ok_cnt}/{len(gds_files)} converted.")
|
||||||
|
|
||||||
|
|
||||||
|
if __name__ == "__main__":
|
||||||
|
main()
|
||||||
68
tools/preview_dataset.py
Normal file
68
tools/preview_dataset.py
Normal file
@@ -0,0 +1,68 @@
|
|||||||
|
#!/usr/bin/env python3
|
||||||
|
"""
|
||||||
|
Quickly preview training pairs (original, transformed, H) from ICLayoutTrainingDataset.
|
||||||
|
Saves a grid image for visual inspection.
|
||||||
|
"""
|
||||||
|
from __future__ import annotations
|
||||||
|
|
||||||
|
import argparse
|
||||||
|
from pathlib import Path
|
||||||
|
|
||||||
|
import numpy as np
|
||||||
|
import torch
|
||||||
|
from PIL import Image
|
||||||
|
from torchvision.utils import make_grid, save_image
|
||||||
|
|
||||||
|
from data.ic_dataset import ICLayoutTrainingDataset
|
||||||
|
from utils.data_utils import get_transform
|
||||||
|
|
||||||
|
|
||||||
|
def to_pil(t: torch.Tensor) -> Image.Image:
|
||||||
|
# input normalized to [-1,1] for 3-channels; invert normalization
|
||||||
|
x = t.clone()
|
||||||
|
if x.dim() == 3 and x.size(0) == 3:
|
||||||
|
x = (x * 0.5) + 0.5 # unnormalize
|
||||||
|
x = (x * 255.0).clamp(0, 255).byte()
|
||||||
|
if x.dim() == 3 and x.size(0) == 3:
|
||||||
|
x = x
|
||||||
|
elif x.dim() == 3 and x.size(0) == 1:
|
||||||
|
x = x.repeat(3, 1, 1)
|
||||||
|
else:
|
||||||
|
raise ValueError("Unexpected tensor shape")
|
||||||
|
np_img = x.permute(1, 2, 0).cpu().numpy()
|
||||||
|
return Image.fromarray(np_img)
|
||||||
|
|
||||||
|
|
||||||
|
def main():
|
||||||
|
parser = argparse.ArgumentParser(description="Preview dataset samples")
|
||||||
|
parser.add_argument("--dir", dest="image_dir", type=str, required=True, help="PNG images directory")
|
||||||
|
parser.add_argument("--out", dest="out_path", type=str, default="preview.png")
|
||||||
|
parser.add_argument("--n", dest="num", type=int, default=8)
|
||||||
|
parser.add_argument("--patch", dest="patch_size", type=int, default=256)
|
||||||
|
parser.add_argument("--elastic", dest="use_elastic", action="store_true")
|
||||||
|
args = parser.parse_args()
|
||||||
|
|
||||||
|
transform = get_transform()
|
||||||
|
ds = ICLayoutTrainingDataset(
|
||||||
|
args.image_dir,
|
||||||
|
patch_size=args.patch_size,
|
||||||
|
transform=transform,
|
||||||
|
scale_range=(1.0, 1.0),
|
||||||
|
use_albu=args.use_elastic,
|
||||||
|
albu_params={"prob": 0.5},
|
||||||
|
)
|
||||||
|
|
||||||
|
images = []
|
||||||
|
for i in range(min(args.num, len(ds))):
|
||||||
|
orig, rot, H = ds[i]
|
||||||
|
# Stack orig and rot side-by-side for each sample
|
||||||
|
images.append(orig)
|
||||||
|
images.append(rot)
|
||||||
|
|
||||||
|
grid = make_grid(torch.stack(images, dim=0), nrow=2, padding=2)
|
||||||
|
save_image(grid, args.out_path)
|
||||||
|
print(f"Saved preview to {args.out_path}")
|
||||||
|
|
||||||
|
|
||||||
|
if __name__ == "__main__":
|
||||||
|
main()
|
||||||
76
tools/smoke_test.py
Normal file
76
tools/smoke_test.py
Normal file
@@ -0,0 +1,76 @@
|
|||||||
|
#!/usr/bin/env python3
|
||||||
|
"""
|
||||||
|
Minimal smoke test:
|
||||||
|
1) Generate a tiny synthetic set (num=8) and rasterize to PNG
|
||||||
|
2) Validate H consistency (n=4, with/without elastic)
|
||||||
|
3) Run a short training loop (epochs=1-2) to verify end-to-end pipeline
|
||||||
|
Prints PASS/FAIL with basic stats.
|
||||||
|
"""
|
||||||
|
from __future__ import annotations
|
||||||
|
|
||||||
|
import argparse
|
||||||
|
import subprocess
|
||||||
|
import os
|
||||||
|
import sys
|
||||||
|
from pathlib import Path
|
||||||
|
|
||||||
|
|
||||||
|
def run(cmd: list[str]) -> int:
|
||||||
|
print("[RUN]", " ".join(cmd))
|
||||||
|
env = os.environ.copy()
|
||||||
|
# Ensure project root on PYTHONPATH for child processes
|
||||||
|
root = Path(__file__).resolve().parents[1]
|
||||||
|
env["PYTHONPATH"] = f"{root}:{env.get('PYTHONPATH','')}" if env.get("PYTHONPATH") else str(root)
|
||||||
|
return subprocess.call(cmd, env=env)
|
||||||
|
|
||||||
|
|
||||||
|
def main() -> None:
|
||||||
|
parser = argparse.ArgumentParser(description="Minimal smoke test for E2E pipeline")
|
||||||
|
parser.add_argument("--root", type=str, default="data/smoke", help="Root dir for smoke test outputs")
|
||||||
|
parser.add_argument("--config", type=str, default="configs/base_config.yaml")
|
||||||
|
args = parser.parse_args()
|
||||||
|
|
||||||
|
root = Path(args.root)
|
||||||
|
gds_dir = root / "gds"
|
||||||
|
png_dir = root / "png"
|
||||||
|
gds_dir.mkdir(parents=True, exist_ok=True)
|
||||||
|
png_dir.mkdir(parents=True, exist_ok=True)
|
||||||
|
|
||||||
|
rc = 0
|
||||||
|
|
||||||
|
# 1) Generate a tiny set
|
||||||
|
rc |= run([sys.executable, "tools/generate_synthetic_layouts.py", "--out_dir", gds_dir.as_posix(), "--num", "8", "--seed", "123"])
|
||||||
|
if rc != 0:
|
||||||
|
print("[FAIL] generate synthetic")
|
||||||
|
sys.exit(2)
|
||||||
|
|
||||||
|
# 2) Rasterize
|
||||||
|
rc |= run([sys.executable, "tools/layout2png.py", "--in", gds_dir.as_posix(), "--out", png_dir.as_posix(), "--dpi", "600"])
|
||||||
|
if rc != 0:
|
||||||
|
print("[FAIL] layout2png")
|
||||||
|
sys.exit(3)
|
||||||
|
|
||||||
|
# 3) Validate H (n=4, both no-elastic and elastic)
|
||||||
|
rc |= run([sys.executable, "tools/validate_h_consistency.py", "--dir", png_dir.as_posix(), "--out", (root/"validate_no_elastic").as_posix(), "--n", "4"])
|
||||||
|
rc |= run([sys.executable, "tools/validate_h_consistency.py", "--dir", png_dir.as_posix(), "--out", (root/"validate_elastic").as_posix(), "--n", "4", "--elastic"])
|
||||||
|
if rc != 0:
|
||||||
|
print("[FAIL] validate H")
|
||||||
|
sys.exit(4)
|
||||||
|
|
||||||
|
# 4) Write back config via synth_pipeline and run short training (1 epoch)
|
||||||
|
rc |= run([sys.executable, "tools/synth_pipeline.py", "--out_root", root.as_posix(), "--num", "0", "--dpi", "600", "--config", args.config, "--ratio", "0.3", "--enable_elastic", "--no_preview"])
|
||||||
|
if rc != 0:
|
||||||
|
print("[FAIL] synth_pipeline config update")
|
||||||
|
sys.exit(5)
|
||||||
|
|
||||||
|
# Train 1 epoch to smoke the loop
|
||||||
|
rc |= run([sys.executable, "train.py", "--config", args.config, "--epochs", "1" ])
|
||||||
|
if rc != 0:
|
||||||
|
print("[FAIL] train 1 epoch")
|
||||||
|
sys.exit(6)
|
||||||
|
|
||||||
|
print("[PASS] Smoke test completed successfully.")
|
||||||
|
|
||||||
|
|
||||||
|
if __name__ == "__main__":
|
||||||
|
main()
|
||||||
169
tools/synth_pipeline.py
Normal file
169
tools/synth_pipeline.py
Normal file
@@ -0,0 +1,169 @@
|
|||||||
|
#!/usr/bin/env python3
|
||||||
|
"""
|
||||||
|
One-click synthetic data pipeline:
|
||||||
|
1) Generate synthetic GDS using tools/generate_synthetic_layouts.py
|
||||||
|
2) Rasterize GDS to PNG using tools/layout2png.py (KLayout preferred, fallback gdstk+SVG)
|
||||||
|
3) Preview random training pairs using tools/preview_dataset.py (optional)
|
||||||
|
4) Validate homography consistency using tools/validate_h_consistency.py (optional)
|
||||||
|
5) Optionally update a YAML config to enable synthetic mixing and elastic augmentation
|
||||||
|
"""
|
||||||
|
from __future__ import annotations
|
||||||
|
|
||||||
|
import argparse
|
||||||
|
import subprocess
|
||||||
|
import sys
|
||||||
|
from pathlib import Path
|
||||||
|
|
||||||
|
from omegaconf import OmegaConf
|
||||||
|
|
||||||
|
|
||||||
|
def run_cmd(cmd: list[str]) -> None:
|
||||||
|
print("[RUN]", " ".join(str(c) for c in cmd))
|
||||||
|
res = subprocess.run(cmd)
|
||||||
|
if res.returncode != 0:
|
||||||
|
raise SystemExit(f"Command failed with code {res.returncode}: {' '.join(map(str, cmd))}")
|
||||||
|
|
||||||
|
|
||||||
|
essential_scripts = {
|
||||||
|
"gen": Path("tools/generate_synthetic_layouts.py"),
|
||||||
|
"gds2png": Path("tools/layout2png.py"),
|
||||||
|
"preview": Path("tools/preview_dataset.py"),
|
||||||
|
"validate": Path("tools/validate_h_consistency.py"),
|
||||||
|
}
|
||||||
|
|
||||||
|
|
||||||
|
def ensure_scripts_exist() -> None:
|
||||||
|
missing = [str(p) for p in essential_scripts.values() if not p.exists()]
|
||||||
|
if missing:
|
||||||
|
raise SystemExit(f"Missing required scripts: {missing}")
|
||||||
|
|
||||||
|
|
||||||
|
def update_config(config_path: Path, png_dir: Path, ratio: float, enable_elastic: bool) -> None:
|
||||||
|
cfg = OmegaConf.load(config_path)
|
||||||
|
# Ensure nodes exist
|
||||||
|
if "synthetic" not in cfg:
|
||||||
|
cfg.synthetic = {}
|
||||||
|
cfg.synthetic.enabled = True
|
||||||
|
cfg.synthetic.png_dir = png_dir.as_posix()
|
||||||
|
cfg.synthetic.ratio = float(ratio)
|
||||||
|
|
||||||
|
if enable_elastic:
|
||||||
|
if "augment" not in cfg:
|
||||||
|
cfg.augment = {}
|
||||||
|
if "elastic" not in cfg.augment:
|
||||||
|
cfg.augment.elastic = {}
|
||||||
|
cfg.augment.elastic.enabled = True
|
||||||
|
# Don't override numeric params if already present
|
||||||
|
if "alpha" not in cfg.augment.elastic:
|
||||||
|
cfg.augment.elastic.alpha = 40
|
||||||
|
if "sigma" not in cfg.augment.elastic:
|
||||||
|
cfg.augment.elastic.sigma = 6
|
||||||
|
if "alpha_affine" not in cfg.augment.elastic:
|
||||||
|
cfg.augment.elastic.alpha_affine = 6
|
||||||
|
if "prob" not in cfg.augment.elastic:
|
||||||
|
cfg.augment.elastic.prob = 0.3
|
||||||
|
# Photometric defaults
|
||||||
|
if "photometric" not in cfg.augment:
|
||||||
|
cfg.augment.photometric = {"brightness_contrast": True, "gauss_noise": True}
|
||||||
|
|
||||||
|
OmegaConf.save(config=cfg, f=config_path)
|
||||||
|
print(f"[OK] Config updated: {config_path}")
|
||||||
|
|
||||||
|
|
||||||
|
def main() -> None:
|
||||||
|
parser = argparse.ArgumentParser(description="One-click synthetic data pipeline")
|
||||||
|
parser.add_argument("--out_root", type=str, default="data/synthetic", help="Root output dir for gds/png/preview")
|
||||||
|
parser.add_argument("--num", type=int, default=200, help="Number of GDS samples to generate")
|
||||||
|
parser.add_argument("--dpi", type=int, default=600, help="Rasterization DPI for PNG rendering")
|
||||||
|
parser.add_argument("--seed", type=int, default=42)
|
||||||
|
parser.add_argument("--ratio", type=float, default=0.3, help="Mixing ratio for synthetic data in training")
|
||||||
|
parser.add_argument("--config", type=str, default="configs/base_config.yaml", help="YAML config to update")
|
||||||
|
parser.add_argument("--enable_elastic", action="store_true", help="Also enable elastic augmentation in config")
|
||||||
|
parser.add_argument("--no_preview", action="store_true", help="Skip preview generation")
|
||||||
|
parser.add_argument("--validate_h", action="store_true", help="Run homography consistency validation on rendered PNGs")
|
||||||
|
parser.add_argument("--validate_n", type=int, default=6, help="Number of samples for H validation")
|
||||||
|
parser.add_argument("--diffusion_dir", type=str, default=None, help="Directory of diffusion-generated PNGs to include")
|
||||||
|
# Rendering style passthrough
|
||||||
|
parser.add_argument("--layermap", type=str, default=None, help="Layer color map for KLayout, e.g. '1/0:#00FF00,2/0:#FF0000'")
|
||||||
|
parser.add_argument("--line_width", type=int, default=None, help="Default draw line width for KLayout display")
|
||||||
|
parser.add_argument("--bgcolor", type=str, default=None, help="Background color for KLayout display")
|
||||||
|
|
||||||
|
args = parser.parse_args()
|
||||||
|
ensure_scripts_exist()
|
||||||
|
|
||||||
|
out_root = Path(args.out_root)
|
||||||
|
gds_dir = out_root / "gds"
|
||||||
|
png_dir = out_root / "png"
|
||||||
|
gds_dir.mkdir(parents=True, exist_ok=True)
|
||||||
|
png_dir.mkdir(parents=True, exist_ok=True)
|
||||||
|
|
||||||
|
# 1) Generate GDS
|
||||||
|
run_cmd([sys.executable, str(essential_scripts["gen"]), "--out_dir", gds_dir.as_posix(), "--num", str(args.num), "--seed", str(args.seed)])
|
||||||
|
|
||||||
|
# 2) GDS -> PNG
|
||||||
|
gds2png_cmd = [
|
||||||
|
sys.executable, str(essential_scripts["gds2png"]),
|
||||||
|
"--in", gds_dir.as_posix(),
|
||||||
|
"--out", png_dir.as_posix(),
|
||||||
|
"--dpi", str(args.dpi),
|
||||||
|
]
|
||||||
|
if args.layermap:
|
||||||
|
gds2png_cmd += ["--layermap", args.layermap]
|
||||||
|
if args.line_width is not None:
|
||||||
|
gds2png_cmd += ["--line_width", str(args.line_width)]
|
||||||
|
if args.bgcolor:
|
||||||
|
gds2png_cmd += ["--bgcolor", args.bgcolor]
|
||||||
|
run_cmd(gds2png_cmd)
|
||||||
|
|
||||||
|
# 3) Preview (optional)
|
||||||
|
if not args.no_preview:
|
||||||
|
preview_path = out_root / "preview.png"
|
||||||
|
preview_cmd = [sys.executable, str(essential_scripts["preview"]), "--dir", png_dir.as_posix(), "--out", preview_path.as_posix(), "--n", "8"]
|
||||||
|
if args.enable_elastic:
|
||||||
|
preview_cmd.append("--elastic")
|
||||||
|
run_cmd(preview_cmd)
|
||||||
|
|
||||||
|
# 4) Validate homography consistency (optional)
|
||||||
|
if args.validate_h:
|
||||||
|
validate_dir = out_root / "validate_h"
|
||||||
|
validate_cmd = [
|
||||||
|
sys.executable, str(essential_scripts["validate"]),
|
||||||
|
"--dir", png_dir.as_posix(),
|
||||||
|
"--out", validate_dir.as_posix(),
|
||||||
|
"--n", str(args.validate_n),
|
||||||
|
]
|
||||||
|
if args.enable_elastic:
|
||||||
|
validate_cmd.append("--elastic")
|
||||||
|
run_cmd(validate_cmd)
|
||||||
|
|
||||||
|
# 5) Update YAML config
|
||||||
|
update_config(Path(args.config), png_dir, args.ratio, args.enable_elastic)
|
||||||
|
# Include diffusion dir if provided (no automatic sampling here; integration only)
|
||||||
|
if args.diffusion_dir:
|
||||||
|
cfg = OmegaConf.load(args.config)
|
||||||
|
if "synthetic" not in cfg:
|
||||||
|
cfg.synthetic = {}
|
||||||
|
if "diffusion" not in cfg.synthetic:
|
||||||
|
cfg.synthetic.diffusion = {}
|
||||||
|
cfg.synthetic.diffusion.enabled = True
|
||||||
|
cfg.synthetic.diffusion.png_dir = Path(args.diffusion_dir).as_posix()
|
||||||
|
# Keep ratio default at 0 unless user updates later; or reuse a small default like 0.1? Keep 0.0 for safety.
|
||||||
|
if "ratio" not in cfg.synthetic.diffusion:
|
||||||
|
cfg.synthetic.diffusion.ratio = 0.0
|
||||||
|
OmegaConf.save(config=cfg, f=args.config)
|
||||||
|
print(f"[OK] Config updated with diffusion_dir: {args.diffusion_dir}")
|
||||||
|
|
||||||
|
print("\n[Done] Synthetic pipeline completed.")
|
||||||
|
print(f"- GDS: {gds_dir}")
|
||||||
|
print(f"- PNG: {png_dir}")
|
||||||
|
if args.diffusion_dir:
|
||||||
|
print(f"- Diffusion PNGs: {Path(args.diffusion_dir)}")
|
||||||
|
if not args.no_preview:
|
||||||
|
print(f"- Preview: {out_root / 'preview.png'}")
|
||||||
|
if args.validate_h:
|
||||||
|
print(f"- H validation: {out_root / 'validate_h'}")
|
||||||
|
print(f"- Updated config: {args.config}")
|
||||||
|
|
||||||
|
|
||||||
|
if __name__ == "__main__":
|
||||||
|
main()
|
||||||
117
tools/validate_h_consistency.py
Normal file
117
tools/validate_h_consistency.py
Normal file
@@ -0,0 +1,117 @@
|
|||||||
|
#!/usr/bin/env python3
|
||||||
|
"""
|
||||||
|
Validate homography consistency produced by ICLayoutTrainingDataset.
|
||||||
|
For random samples, we check that cv2.warpPerspective(original, H) ≈ transformed.
|
||||||
|
Saves visual composites and prints basic metrics (MSE / PSNR).
|
||||||
|
"""
|
||||||
|
from __future__ import annotations
|
||||||
|
|
||||||
|
import argparse
|
||||||
|
from pathlib import Path
|
||||||
|
import sys
|
||||||
|
|
||||||
|
import cv2
|
||||||
|
import numpy as np
|
||||||
|
import torch
|
||||||
|
from PIL import Image
|
||||||
|
|
||||||
|
# Ensure project root is on sys.path when running as a script
|
||||||
|
PROJECT_ROOT = Path(__file__).resolve().parents[1]
|
||||||
|
if str(PROJECT_ROOT) not in sys.path:
|
||||||
|
sys.path.insert(0, str(PROJECT_ROOT))
|
||||||
|
|
||||||
|
from data.ic_dataset import ICLayoutTrainingDataset
|
||||||
|
|
||||||
|
|
||||||
|
def tensor_to_u8_img(t: torch.Tensor) -> np.ndarray:
|
||||||
|
"""Convert 1xHxW or 3xHxW float tensor in [0,1] to uint8 HxW or HxWx3."""
|
||||||
|
if t.dim() != 3:
|
||||||
|
raise ValueError(f"Expect 3D tensor, got {t.shape}")
|
||||||
|
if t.size(0) == 1:
|
||||||
|
arr = (t.squeeze(0).cpu().numpy() * 255.0).clip(0, 255).astype(np.uint8)
|
||||||
|
elif t.size(0) == 3:
|
||||||
|
arr = (t.permute(1, 2, 0).cpu().numpy() * 255.0).clip(0, 255).astype(np.uint8)
|
||||||
|
else:
|
||||||
|
raise ValueError(f"Unexpected channels: {t.size(0)}")
|
||||||
|
return arr
|
||||||
|
|
||||||
|
|
||||||
|
def mse(a: np.ndarray, b: np.ndarray) -> float:
|
||||||
|
diff = a.astype(np.float32) - b.astype(np.float32)
|
||||||
|
return float(np.mean(diff * diff))
|
||||||
|
|
||||||
|
|
||||||
|
def psnr(a: np.ndarray, b: np.ndarray) -> float:
|
||||||
|
m = mse(a, b)
|
||||||
|
if m <= 1e-8:
|
||||||
|
return float('inf')
|
||||||
|
return 10.0 * np.log10((255.0 * 255.0) / m)
|
||||||
|
|
||||||
|
|
||||||
|
def main() -> None:
|
||||||
|
parser = argparse.ArgumentParser(description="Validate homography consistency")
|
||||||
|
parser.add_argument("--dir", dest="image_dir", type=str, required=True, help="PNG images directory")
|
||||||
|
parser.add_argument("--out", dest="out_dir", type=str, default="validate_h_out", help="Output directory for composites")
|
||||||
|
parser.add_argument("--n", dest="num", type=int, default=8, help="Number of samples to validate")
|
||||||
|
parser.add_argument("--patch", dest="patch_size", type=int, default=256)
|
||||||
|
parser.add_argument("--elastic", dest="use_elastic", action="store_true")
|
||||||
|
args = parser.parse_args()
|
||||||
|
|
||||||
|
out_dir = Path(args.out_dir)
|
||||||
|
out_dir.mkdir(parents=True, exist_ok=True)
|
||||||
|
|
||||||
|
# Use no photometric/Sobel transform here to compare raw grayscale content
|
||||||
|
ds = ICLayoutTrainingDataset(
|
||||||
|
args.image_dir,
|
||||||
|
patch_size=args.patch_size,
|
||||||
|
transform=None,
|
||||||
|
scale_range=(1.0, 1.0),
|
||||||
|
use_albu=args.use_elastic,
|
||||||
|
albu_params={"prob": 0.5},
|
||||||
|
)
|
||||||
|
|
||||||
|
n = min(args.num, len(ds))
|
||||||
|
if n == 0:
|
||||||
|
print("[WARN] Empty dataset.")
|
||||||
|
return
|
||||||
|
|
||||||
|
mses = []
|
||||||
|
psnrs = []
|
||||||
|
|
||||||
|
for i in range(n):
|
||||||
|
patch_t, trans_t, H2x3_t = ds[i]
|
||||||
|
# Convert to uint8 arrays
|
||||||
|
patch_u8 = tensor_to_u8_img(patch_t)
|
||||||
|
trans_u8 = tensor_to_u8_img(trans_t)
|
||||||
|
if patch_u8.ndim == 3:
|
||||||
|
patch_u8 = cv2.cvtColor(patch_u8, cv2.COLOR_BGR2GRAY)
|
||||||
|
if trans_u8.ndim == 3:
|
||||||
|
trans_u8 = cv2.cvtColor(trans_u8, cv2.COLOR_BGR2GRAY)
|
||||||
|
|
||||||
|
# Reconstruct 3x3 H
|
||||||
|
H2x3 = H2x3_t.numpy()
|
||||||
|
H = np.vstack([H2x3, [0.0, 0.0, 1.0]]).astype(np.float32)
|
||||||
|
|
||||||
|
# Warp original with H
|
||||||
|
warped = cv2.warpPerspective(patch_u8, H, (patch_u8.shape[1], patch_u8.shape[0]))
|
||||||
|
|
||||||
|
# Metrics
|
||||||
|
m = mse(warped, trans_u8)
|
||||||
|
p = psnr(warped, trans_u8)
|
||||||
|
mses.append(m)
|
||||||
|
psnrs.append(p)
|
||||||
|
|
||||||
|
# Composite image: [orig | warped | transformed | absdiff]
|
||||||
|
diff = cv2.absdiff(warped, trans_u8)
|
||||||
|
comp = np.concatenate([
|
||||||
|
patch_u8, warped, trans_u8, diff
|
||||||
|
], axis=1)
|
||||||
|
out_path = out_dir / f"sample_{i:03d}.png"
|
||||||
|
cv2.imwrite(out_path.as_posix(), comp)
|
||||||
|
print(f"[OK] sample {i}: MSE={m:.2f}, PSNR={p:.2f} dB -> {out_path}")
|
||||||
|
|
||||||
|
print(f"\nSummary: MSE avg={np.mean(mses):.2f} ± {np.std(mses):.2f}, PSNR avg={np.mean(psnrs):.2f} dB")
|
||||||
|
|
||||||
|
|
||||||
|
if __name__ == "__main__":
|
||||||
|
main()
|
||||||
147
train.py
147
train.py
@@ -7,7 +7,7 @@ from datetime import datetime
|
|||||||
from pathlib import Path
|
from pathlib import Path
|
||||||
|
|
||||||
import torch
|
import torch
|
||||||
from torch.utils.data import DataLoader
|
from torch.utils.data import DataLoader, ConcatDataset, WeightedRandomSampler
|
||||||
from torch.utils.tensorboard import SummaryWriter
|
from torch.utils.tensorboard import SummaryWriter
|
||||||
|
|
||||||
from data.ic_dataset import ICLayoutTrainingDataset
|
from data.ic_dataset import ICLayoutTrainingDataset
|
||||||
@@ -82,25 +82,152 @@ def main(args):
|
|||||||
|
|
||||||
transform = get_transform()
|
transform = get_transform()
|
||||||
|
|
||||||
dataset = ICLayoutTrainingDataset(
|
# 读取增强与合成配置
|
||||||
|
augment_cfg = cfg.get("augment", {})
|
||||||
|
elastic_cfg = augment_cfg.get("elastic", {}) if augment_cfg else {}
|
||||||
|
use_albu = bool(elastic_cfg.get("enabled", False))
|
||||||
|
albu_params = {
|
||||||
|
"prob": elastic_cfg.get("prob", 0.3),
|
||||||
|
"alpha": elastic_cfg.get("alpha", 40),
|
||||||
|
"sigma": elastic_cfg.get("sigma", 6),
|
||||||
|
"alpha_affine": elastic_cfg.get("alpha_affine", 6),
|
||||||
|
"brightness_contrast": bool(augment_cfg.get("photometric", {}).get("brightness_contrast", True)) if augment_cfg else True,
|
||||||
|
"gauss_noise": bool(augment_cfg.get("photometric", {}).get("gauss_noise", True)) if augment_cfg else True,
|
||||||
|
}
|
||||||
|
|
||||||
|
# 构建真实数据集
|
||||||
|
real_dataset = ICLayoutTrainingDataset(
|
||||||
data_dir,
|
data_dir,
|
||||||
patch_size=patch_size,
|
patch_size=patch_size,
|
||||||
transform=transform,
|
transform=transform,
|
||||||
scale_range=scale_range,
|
scale_range=scale_range,
|
||||||
|
use_albu=use_albu,
|
||||||
|
albu_params=albu_params,
|
||||||
)
|
)
|
||||||
|
|
||||||
logger.info(f"数据集大小: {len(dataset)}")
|
# 读取合成数据配置（程序化 + 扩散）
|
||||||
|
syn_cfg = cfg.get("synthetic", {})
|
||||||
|
syn_enabled = bool(syn_cfg.get("enabled", False))
|
||||||
|
syn_ratio = float(syn_cfg.get("ratio", 0.0))
|
||||||
|
syn_dir = syn_cfg.get("png_dir", None)
|
||||||
|
|
||||||
# 分割训练集和验证集
|
syn_dataset = None
|
||||||
train_size = int(0.8 * len(dataset))
|
if syn_enabled and syn_dir:
|
||||||
val_size = len(dataset) - train_size
|
syn_dir_path = Path(to_absolute_path(syn_dir, config_dir))
|
||||||
train_dataset, val_dataset = torch.utils.data.random_split(dataset, [train_size, val_size])
|
if syn_dir_path.exists():
|
||||||
|
syn_dataset = ICLayoutTrainingDataset(
|
||||||
|
syn_dir_path.as_posix(),
|
||||||
|
patch_size=patch_size,
|
||||||
|
transform=transform,
|
||||||
|
scale_range=scale_range,
|
||||||
|
use_albu=use_albu,
|
||||||
|
albu_params=albu_params,
|
||||||
|
)
|
||||||
|
if len(syn_dataset) == 0:
|
||||||
|
syn_dataset = None
|
||||||
|
else:
|
||||||
|
logger.warning(f"合成数据目录不存在，忽略: {syn_dir_path}")
|
||||||
|
syn_enabled = False
|
||||||
|
|
||||||
logger.info(f"训练集大小: {len(train_dataset)}, 验证集大小: {len(val_dataset)}")
|
# 扩散生成数据配置
|
||||||
|
diff_cfg = syn_cfg.get("diffusion", {}) if syn_cfg else {}
|
||||||
|
diff_enabled = bool(diff_cfg.get("enabled", False))
|
||||||
|
diff_ratio = float(diff_cfg.get("ratio", 0.0))
|
||||||
|
diff_dir = diff_cfg.get("png_dir", None)
|
||||||
|
diff_dataset = None
|
||||||
|
if diff_enabled and diff_dir:
|
||||||
|
diff_dir_path = Path(to_absolute_path(diff_dir, config_dir))
|
||||||
|
if diff_dir_path.exists():
|
||||||
|
diff_dataset = ICLayoutTrainingDataset(
|
||||||
|
diff_dir_path.as_posix(),
|
||||||
|
patch_size=patch_size,
|
||||||
|
transform=transform,
|
||||||
|
scale_range=scale_range,
|
||||||
|
use_albu=use_albu,
|
||||||
|
albu_params=albu_params,
|
||||||
|
)
|
||||||
|
if len(diff_dataset) == 0:
|
||||||
|
diff_dataset = None
|
||||||
|
else:
|
||||||
|
logger.warning(f"扩散数据目录不存在，忽略: {diff_dir_path}")
|
||||||
|
diff_enabled = False
|
||||||
|
|
||||||
|
logger.info(
|
||||||
|
"真实数据集大小: %d%s%s" % (
|
||||||
|
len(real_dataset),
|
||||||
|
f", 合成(程序)数据集: {len(syn_dataset)}" if syn_dataset else "",
|
||||||
|
f", 合成(扩散)数据集: {len(diff_dataset)}" if diff_dataset else "",
|
||||||
|
)
|
||||||
|
)
|
||||||
|
|
||||||
|
# 验证集仅使用真实数据，避免评价受合成样本干扰
|
||||||
|
train_size = int(0.8 * len(real_dataset))
|
||||||
|
val_size = max(len(real_dataset) - train_size, 1)
|
||||||
|
real_train_dataset, val_dataset = torch.utils.data.random_split(real_dataset, [train_size, val_size])
|
||||||
|
|
||||||
|
# 训练集：可与合成数据集合并（程序合成 + 扩散）
|
||||||
|
datasets = [real_train_dataset]
|
||||||
|
weights = []
|
||||||
|
names = []
|
||||||
|
# 收集各源与期望比例
|
||||||
|
n_real = len(real_train_dataset)
|
||||||
|
n_real = max(n_real, 1)
|
||||||
|
names.append("real")
|
||||||
|
# 程序合成
|
||||||
|
if syn_dataset is not None and syn_enabled and syn_ratio > 0.0:
|
||||||
|
datasets.append(syn_dataset)
|
||||||
|
names.append("synthetic")
|
||||||
|
# 扩散合成
|
||||||
|
if diff_dataset is not None and diff_enabled and diff_ratio > 0.0:
|
||||||
|
datasets.append(diff_dataset)
|
||||||
|
names.append("diffusion")
|
||||||
|
|
||||||
|
if len(datasets) > 1:
|
||||||
|
mixed_train_dataset = ConcatDataset(datasets)
|
||||||
|
# 计算各源样本数
|
||||||
|
counts = [len(real_train_dataset)]
|
||||||
|
if syn_dataset is not None and syn_enabled and syn_ratio > 0.0:
|
||||||
|
counts.append(len(syn_dataset))
|
||||||
|
if diff_dataset is not None and diff_enabled and diff_ratio > 0.0:
|
||||||
|
counts.append(len(diff_dataset))
|
||||||
|
# 期望比例：real = 1 - (syn_ratio + diff_ratio)
|
||||||
|
target_real = max(0.0, 1.0 - (syn_ratio + diff_ratio))
|
||||||
|
target_ratios = [target_real]
|
||||||
|
if syn_dataset is not None and syn_enabled and syn_ratio > 0.0:
|
||||||
|
target_ratios.append(syn_ratio)
|
||||||
|
if diff_dataset is not None and diff_enabled and diff_ratio > 0.0:
|
||||||
|
target_ratios.append(diff_ratio)
|
||||||
|
# 构建每个样本的权重
|
||||||
|
per_source_weights = []
|
||||||
|
for count, ratio in zip(counts, target_ratios):
|
||||||
|
count = max(count, 1)
|
||||||
|
per_source_weights.append(ratio / count)
|
||||||
|
# 展开到每个样本
|
||||||
|
weights = []
|
||||||
|
idx = 0
|
||||||
|
for count, w in zip(counts, per_source_weights):
|
||||||
|
weights += [w] * count
|
||||||
|
idx += count
|
||||||
|
sampler = WeightedRandomSampler(weights, num_samples=len(mixed_train_dataset), replacement=True)
|
||||||
|
train_dataloader = DataLoader(mixed_train_dataset, batch_size=batch_size, sampler=sampler, num_workers=4)
|
||||||
|
logger.info(
|
||||||
|
f"启用混采: real={target_real:.2f}, syn={syn_ratio:.2f}, diff={diff_ratio:.2f}; 总样本={len(mixed_train_dataset)}"
|
||||||
|
)
|
||||||
if writer:
|
if writer:
|
||||||
writer.add_text("dataset/info", f"train={len(train_dataset)}, val={len(val_dataset)}")
|
writer.add_text(
|
||||||
|
"dataset/mix",
|
||||||
|
f"enabled=true, ratios: real={target_real:.2f}, syn={syn_ratio:.2f}, diff={diff_ratio:.2f}; "
|
||||||
|
f"counts: real_train={len(real_train_dataset)}, syn={len(syn_dataset) if syn_dataset else 0}, diff={len(diff_dataset) if diff_dataset else 0}"
|
||||||
|
)
|
||||||
|
else:
|
||||||
|
train_dataloader = DataLoader(real_train_dataset, batch_size=batch_size, shuffle=True, num_workers=4)
|
||||||
|
if writer:
|
||||||
|
writer.add_text("dataset/mix", f"enabled=false, real_train={len(real_train_dataset)}")
|
||||||
|
|
||||||
|
logger.info(f"训练集大小: {len(train_dataloader.dataset)}, 验证集大小: {len(val_dataset)}")
|
||||||
|
if writer:
|
||||||
|
writer.add_text("dataset/info", f"train={len(train_dataloader.dataset)}, val={len(val_dataset)}")
|
||||||
|
|
||||||
train_dataloader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True, num_workers=4)
|
|
||||||
val_dataloader = DataLoader(val_dataset, batch_size=batch_size, shuffle=False, num_workers=4)
|
val_dataloader = DataLoader(val_dataset, batch_size=batch_size, shuffle=False, num_workers=4)
|
||||||
|
|
||||||
model = RoRD().cuda()
|
model = RoRD().cuda()
|
||||||
|
|||||||
Reference in New Issue
Block a user