From cdd8c43e2b851a5f44035b12c20fa249a8471ddb Mon Sep 17 00:00:00 2001 From: hxh5159 Date: Sun, 31 May 2026 09:45:08 +0800 Subject: [PATCH] 001 --- ASSIGNMENT_GUIDE.md | 454 ++++++++++++++++++++ hw_monitor.py | 264 ++++++++++++ inference.py | 302 +++++++++++++ run_experiments.sh | 188 ++++++++ training/config/ablation_no_camera.yaml | 157 +++++++ training/config/ablation_no_depth.yaml | 158 +++++++ training/config/ablation_no_grad_depth.yaml | 170 ++++++++ training/config/full_finetune.yaml | 183 ++++++++ training/launch.py | 2 - training/launch_multi.py | 119 +++++ training/loss.py | 6 +- training/trainer.py | 1 - 12 files changed, 1998 insertions(+), 6 deletions(-) create mode 100644 ASSIGNMENT_GUIDE.md create mode 100644 hw_monitor.py create mode 100644 inference.py create mode 100644 run_experiments.sh create mode 100644 training/config/ablation_no_camera.yaml create mode 100644 training/config/ablation_no_depth.yaml create mode 100644 training/config/ablation_no_grad_depth.yaml create mode 100644 training/config/full_finetune.yaml create mode 100644 training/launch_multi.py diff --git a/ASSIGNMENT_GUIDE.md b/ASSIGNMENT_GUIDE.md new file mode 100644 index 000000000..9ed7b4fe2 --- /dev/null +++ b/ASSIGNMENT_GUIDE.md @@ -0,0 +1,454 @@ +# VGGT CV课程作业完整指南 + +## 项目概述 + +VGGT (Visual Geometry Grounded Transformer) 是 Meta (Facebook Research) 提出的一个从图像序列直接预测 3D 几何信息的端到端模型，能够同时估计相机位姿、深度图和 3D 点云。 + +## 架构理解 + +### 模型整体架构 + +``` +输入图像 [B, S, 3, H, W] + │ + ▼ +┌─────────────────────────────────────┐ +│ Aggregator (ViT) │ ← 核心: 交错注意力机制 +│ ┌─────────────────────────────┐ │ +│ │ Frame Attention (空间) │ │ 对每帧内部做 self-attention +│ │ Global Attention (时空) │ │ 对所有帧的 tokens 做全局 attention +│ └─────────────────────────────┘ │ +│ depth=24, embed_dim=1024 │ +│ patch_size=14 │ +│ 使用 DINOv2 ViT-L 预训练权重 │ +└─────────────────────────────────────┘ + │ aggregated_tokens_list (24 层输出) + ├──────────────────┬──────────────────┐ + ▼ ▼ ▼ +┌──────────────┐ ┌──────────────┐ ┌──────────────┐ +│ CameraHead │ │ DPTHead │ │ TrackHead │ +│ │ │ (Depth) │ │ (Tracking) │ +│ 4次迭代优化 │ │ │ │ │ +│ 预测: │ │ 预测: │ │ 预测: │ +│ - T (平移) │ │ - depth │ │ - tracks │ +│ - R (旋转) │ │ - depth_conf │ │ - vis │ +│ - FL (焦距) │ │ │ │ - conf │ +└──────────────┘ └──────────────┘ └──────────────┘ +``` + +### 关键设计特点 +- **Alternating Attention**: 在 Frame Attention (空间) 和 Global Attention (时空) 之间交替 +- **Special Tokens**: Camera Token 和 Register Token 用于捕获全局信息 +- **Iterative Refinement**: CameraHead 使用 4 次迭代从 coarse-to-fine 优化位姿 +- **DPT Head**: 使用类似 DepthAnything V2 的多尺度特征融合架构 +- **Confidence Prediction**: 深度和点云预测都带有置信度，使用 `gamma * loss * conf - alpha * log(conf)` 损失 + +### Loss 函数详细分析 + +``` +总损失 = 5.0 × Camera Loss + 1.0 × Depth Loss + +Camera Loss 包含: + ├── loss_T (Translation): |pred_T - gt_T| (L1) + ├── loss_R (Rotation): |pred_quat - gt_quat| (L1) + └── loss_FL (Focal Length): |pred_FL - gt_FL| (L1) + +Depth Loss 包含: + ├── loss_reg_depth: ||pred_depth - gt_depth||² (基础回归损失) + ├── loss_conf_depth: γ × loss_reg × conf - α × log(conf) (置信度加权) + └── loss_grad_depth: |grad(pred) - grad(gt)| (梯度平滑损失) + └── 多尺度 (3层): 1x, 2x, 4x 下采样 +``` + +### 当前训练配置 (default.yaml) +- **部分微调**: 冻结 Aggregator (`"*aggregator*"`)，只训练 CameraHead 和 DPTHead +- **Co3D 数据**: 只使用 `apple` 类别 (debug=True) +- **图像**: 518×518, patch_size=14 +- **AMP**: bfloat16 混合精度 +- **梯度累积**: 3 步 + +--- + +## 任务 1: 全参微调 + +### 修改内容 + +#### 新增文件: `training/config/full_finetune.yaml` + +与默认配置的关键区别: +```yaml +# 移除了冻结设置，允许所有参数训练 +frozen_module_names: [] # 原来是 ["*aggregator*"] + +# 增大梯度累积步数以补偿单卡更小的 batch +accum_steps: 4 # 原来是 3 +``` + +### 修改文件: `training/trainer.py` (L540) +- **修复**: 移除了验证循环中的 `import pdb; pdb.set_trace()` 调试断点 + +### 修改文件: `training/launch.py` +- **修复**: 移除了末尾的 `import pdb; pdb.set_trace()` 调试代码 + +### 运行命令 +```bash +# 全参微调 (需要 H100 80GB) +torchrun --nproc_per_node=1 training/launch_multi.py --config full_finetune \ + --override \ + "data.train.dataset.dataset_configs.0.CO3D_DIR=/your/path/to/co3d" \ + "data.train.dataset.dataset_configs.0.CO3D_ANNOTATION_DIR=/your/path/to/co3d_anno" \ + "data.val.dataset.dataset_configs.0.CO3D_DIR=/your/path/to/co3d" \ + "data.val.dataset.dataset_configs.0.CO3D_ANNOTATION_DIR=/your/path/to/co3d_anno" +``` + +### 为什么 4090 无法支持全参微调? + +| 组件 | 参数量 | 备注 | +|------|--------|------| +| Aggregator (ViT-L DINOv2) | ~300M | 24 层 Transformer, embed_dim=1024 | +| CameraHead | ~50M | 4 层 trunk + pose branch | +| DPTHead (Depth) | ~50M | 多尺度融合 + 上采样 | +| **总参数量** | **~400M** | | +| 全参微调 (bf16 模型) | ~0.8 GB | | +| AdamW 优化器状态 | ~3.2 GB | 2× 模型参数 (momentum + variance) | +| 梯度 | ~0.8 GB | | +| 激活值 (batch=1, seq=12, 518²) | ~50-60 GB | **主要的显存消耗** | +| **总显存需求** | **~55-70 GB** | | +| 4090 (24GB) | ❌ 不足以容纳 | | +| H100 (80GB) | ✅ 可以运行 | | + +--- + +## 任务 2: Toy 消融实验 - 哪个 Loss 可以在微调阶段去掉 + +### 实验设计 + +我设计了 4 组对比实验: + +| 实验 | 配置文件 | Camera Loss | Depth Reg Loss | Depth Conf Loss | Depth Grad Loss | +|------|----------|:-----------:|:--------------:|:---------------:|:---------------:| +| 基线 | `default.yaml` | ✅ (w=5.0) | ✅ | ✅ | ✅ (grad) | +| 消融1 | `ablation_no_depth.yaml` | ✅ (w=5.0) | ❌ | ❌ | ❌ | +| 消融2 | `ablation_no_camera.yaml` | ❌ | ✅ | ✅ | ✅ (grad) | +| 消融3 | `ablation_no_grad_depth.yaml` | ✅ (w=5.0) | ✅ | ✅ | ❌ (null) | + +### 新增配置文件 +- `training/config/ablation_no_depth.yaml` — 仅 Camera Loss +- `training/config/ablation_no_camera.yaml` — 仅 Depth Loss +- `training/config/ablation_no_grad_depth.yaml` — Camera + Depth 但无梯度损失 + +### 代码修改: `training/loss.py` +在 `regression_loss` 函数中增加对 `gradient_loss_fn=None` 的处理: +```python +# 修改前: +if "conf" in gradient_loss_fn: # None 会导致 TypeError +if "normal" in gradient_loss_fn: + +# 修改后: +if gradient_loss_fn is not None and "conf" in gradient_loss_fn: +if gradient_loss_fn is not None and "normal" in gradient_loss_fn: +elif gradient_loss_fn is not None and "grad" in gradient_loss_fn: +``` + +### 运行命令 +```bash +# 消融实验 1: 无深度损失 +torchrun --nproc_per_node=1 training/launch_multi.py --config ablation_no_depth + +# 消融实验 2: 无相机损失 +torchrun --nproc_per_node=1 training/launch_multi.py --config ablation_no_camera + +# 消融实验 3: 无深度梯度损失 +torchrun --nproc_per_node=1 training/launch_multi.py --config ablation_no_grad_depth +``` + +### 预期结论 + +根据架构分析，我的假设是: + +1. **Camera Loss (loss_camera) 不能去掉** — 这是核心任务，去掉后模型无法学习正确的相机位姿，所有下游任务都会受影响。权重为 5.0 也说明了它的重要性。 + +2. **Depth Loss (loss_reg_depth + loss_conf_depth) 可以在微调阶段去掉** — 在 Co3D 数据集上微调时，如果主要目标是提升相机位姿估计精度，深度损失作为辅助任务可能不是必需的。预训练模型已经学到了良好的深度先验。 + +3. **Depth Gradient Loss (loss_grad_depth) 是最可去掉的** — 它只是一个空间平滑正则化项。在预训练已经足够好的情况下，微调阶段去掉这个损失影响最小。这个损失占据了额外的计算（多尺度梯度计算）但对核心精度贡献有限。 + +**推荐结论**: `loss_grad_depth` (深度梯度损失) 是微调阶段最可以去掉的 loss function。 + +--- + +## 任务 3: VGGT 的局限性与未来改进方向 + +### 当前局限性 + +1. **固定分辨率限制**: VGGT 固定在 518×518 分辨率，对高分辨率细节捕捉不足 + - 改进方向: 多尺度输入、自适应分辨率处理 + +2. **单目深度估计的不确定性**: 深度预测是单目的，缺乏多视图几何约束 + - 改进方向: 结合多视图立体匹配 (MVS) 约束 + +3. **Tracking 模块未充分集成**: 当前代码中 track loss 是注释掉的 (dirty code) + - 改进方向: 完善 tracking 模块，实现端到端的 tracking + 3D 联合训练 + +4. **类别泛化有限**: Co3D 只包含 51 个物体类别，对场景级别数据 (如室内/室外场景) 泛化能力未知 + - 改进方向: 在大规模场景数据 (如 ScanNet, MegaDepth) 上进行混合训练 + +5. **序列长度限制**: 长序列 (>24 帧) 会导致 Global Attention 的 O(S²) 复杂度问题 + - 改进方向: 使用滑动窗口 attention 或 linear attention + +6. **缺乏不确定性量化**: 虽然深度有置信度预测，但相机位姿没有不确定性估计 + - 改进方向: 添加位姿的协方差预测 (probabilistic pose estimation) + +7. **不支持动态场景**: 假设场景是静态的，无法处理运动物体 + - 改进方向: 添加运动分割模块，对动态物体单独建模 + +### 如果让我 Follow VGGT，我会添加的新能力 + +1. **多模态输入融合**: 结合语义分割/目标检测信息，增强对场景的理解 + - 例如: 使用 SAM (Segment Anything Model) 的 mask 作为额外条件 + +2. **在线学习能力**: 支持在新场景上快速自适应 (test-time adaptation) + - 通过最小化光度重投影误差在线微调 + +3. **层次化场景表示**: 同时输出稀疏 (关键点) 和稠密 (per-pixel) 的 3D 表示 + - 类似 Gaussian Splatting 的混合表示 + +4. **闭环检测与全局一致优化**: 当前每帧独立预测，缺少全局一致性 + - 添加可微的 BA (Bundle Adjustment) 层，端到端优化全局一致性 + +5. **时序平滑先验**: 添加时序一致性损失，使相邻帧的预测更加平滑 + - 使用 temporal smoothness loss 约束相邻帧的位姿和深度变化 + +--- + +## 任务 4: TensorBoard 各项指标详解 + +### 训练指标 (Train) + +| 指标名称 | 含义 | 期望趋势 | +|----------|------|----------| +| `Loss/train_loss_objective` | **总训练损失** (所有加权子损失之和) | ↓ 下降 | +| `Loss/train_loss_camera` | **相机位姿损失** (T + R + FL 加权和) | ↓ 下降 | +| `Loss/train_loss_T` | **平移损失** \|pred_T - gt_T\| (L1) | ↓ 下降 | +| `Loss/train_loss_R` | **旋转损失** \|pred_quat - gt_quat\| (L1) | ↓ 下降 | +| `Loss/train_loss_FL` | **焦距损失** \|pred_FL - gt_FL\| (L1) | ↓ 下降 | +| `Loss/train_loss_conf_depth` | **深度置信度加权损失** γ·reg·conf - α·log(conf) | ↓ 下降 (但不会到0) | +| `Loss/train_loss_reg_depth` | **深度回归损失** \|pred_depth - gt_depth\|² | ↓ 下降 | +| `Loss/train_loss_grad_depth` | **深度梯度损失** 空间平滑约束 (多尺度) | ↓ 下降 | + +### 梯度指标 (Grad) + +| 指标名称 | 含义 | +|----------|------| +| `Grad/aggregator` | Aggregator 模块的梯度范数 (被裁剪前) | +| `Grad/depth` | Depth Head 模块的梯度范数 | +| `Grad/camera` | Camera Head 模块的梯度范数 | + +如果梯度范数突然变得很大 → 可能出现训练不稳定 +如果梯度范数接近0 → 可能是梯度消失 + +### 优化器指标 (Optim) + +| 指标名称 | 含义 | +|----------|------| +| `Optim/lr` | 当前学习率 (warmup → cosine decay) | +| `Optim/weight_decay` | 当前权重衰减率 | +| `Optim/where` | 训练进度 [0, 1]，0=开始，1=结束 | + +### 验证指标 (Val) + +| 指标名称 | 含义 | +|----------|------| +| `Loss/val_loss_objective` | 验证集总损失 | +| `Loss/val_loss_camera` | 验证集相机位姿损失 | +| `Loss/val_loss_T` | 验证集平移损失 | +| `Loss/val_loss_R` | 验证集旋转损失 | +| `Loss/val_loss_FL` | 验证集焦距损失 | +| `Loss/val_loss_conf_depth` | 验证集深度置信度损失 | +| `Loss/val_loss_reg_depth` | 验证集深度回归损失 | +| `Loss/val_loss_grad_depth` | 验证集深度梯度损失 | +| `Trainer/where` | 当前训练进度 | +| `Trainer/epoch` | 当前 epoch 数 | +| `Trainer/steps_val` | 验证步数 | + +### 如何解读 TensorBoard + +1. **正常训练**: 所有 Loss 曲线平稳下降，Grad 范数稳定在 0.1-10 范围内 +2. **过拟合**: Train loss 继续下降但 Val loss 开始上升 +3. **梯度爆炸**: Grad 范数突然跳到 100+，对应 Loss 出现 spike +4. **学习率过大**: Loss 震荡不收敛，考虑降低初始 lr +5. **某个 loss 不下降**: 检查该 loss 的权重是否太小，或数据是否有问题 + +--- + +## 任务 5: 样本外数据推理 + +### 新增文件: `inference.py` + +功能: +- 加载预训练/微调后的 VGGT 模型 +- 对任意文件夹中的图像进行 3D 重建 +- 输出相机位姿 (cameras.json)、深度图 (depth_maps.npy)、3D 点云 (points_3d.npy) + +### 运行命令 + +```bash +# 使用 HuggingFace 预训练模型推理 +python inference.py \ + --image_dir /path/to/your/images \ + --output_dir ./vggt_output + +# 使用微调后的模型推理 +python inference.py \ + --image_dir /path/to/your/images \ + --output_dir ./vggt_output \ + --checkpoint logs/exp001/ckpts/checkpoint.pt +``` + +### 输入要求 +- 图像目录包含 ≥2 张同一场景的不同视角图像 +- 支持 jpg, jpeg, png 格式 +- 建议图像分辨率 ≥1024px，清晰无模糊 + +### 输出说明 +- `cameras.json`: 每张图像的相机外参 (R|t) 和内参 (K 矩阵) +- `depth_maps.npy`: 预测的深度图 (H×W, 米为单位) +- `depth_confs.npy`: 深度置信度 (H×W, 越高越可信) +- `depth_stats.json`: 深度统计信息 +- `points_3d.npy`: 高置信度区域的 3D 点云 + +--- + +## 任务 6: 硬件使用记录 + +### 新增文件: `hw_monitor.py` + +功能: +- 启动后台线程定时采样 GPU 和 CPU 状态 (通过 nvidia-smi) +- 记录 PyTorch 显存分配状态 +- 保存原始数据为 JSON，生成人类可读的报告 + +### 运行命令 + +```bash +# 终端 1: 启动硬件监控（后台运行） +python hw_monitor.py --interval 2 --output hw_stats.json + +# 终端 2: 运行训练 +torchrun --nproc_per_node=1 training/launch_multi.py --config default + +# 训练结束后，在终端1按 Ctrl+C 停止监控 +# 会自动打印硬件使用报告 + +# 查看已有报告的摘要 +python hw_monitor.py --report hw_stats.json +``` + +### 预期硬件使用对比 + +| 实验 | GPU | 显存峰值 | GPU利用率 | 原因分析 | +|------|-----|----------|-----------|----------| +| 部分微调 (default) | 4090 24GB | ~18-22 GB | 80-95% | 只训练 head，Aggregator 冻结 | +| 部分微调 (default) | H100 80GB | ~18-22 GB | 60-80% | 同上，剩余显存空闲 | +| 全参微调 (full_finetune) | H100 80GB | ~60-70 GB | 90-100% | 所有参数可训练，激活值占用大 | +| 推理 | 任意 | ~8-12 GB | 50-80% | 无反向传播，显存需求显著降低 | + +### 与上次作业的差异分析 + +VGGT 与一般 CV 模型（如 ResNet/分类模型）的硬件使用差异: + +1. **显存消耗更大**: VGGT 使用 ViT-L 骨干 (~300M 参数) + 多个预测头 + 全图激活值，显存是典型分类模型的 5-10 倍 +2. **序列处理**: 同时处理多帧图像 (2-24 帧)，激活值与帧数线性增长 +3. **AMP (bfloat16) 关键**: 不使用混合精度训练时，显存需求翻倍 +4. **Gradient Checkpointing**: Aggregator 的 frame/global blocks 使用 `torch.utils.checkpoint` 来减少激活值显存 + +--- + +## 项目修改总结 + +### 修改的文件 + +| 文件 | 修改内容 | 原因 | +|------|----------|------| +| `training/trainer.py:L540` | 移除 `import pdb; pdb.set_trace()` | 修复调试断点，否则验证时会停止 | +| `training/launch.py:L9-11` | 移除 `import pdb; pdb.set_trace(); m=1` | 修复调试代码 | +| `training/loss.py:L319-336` | 增加 `gradient_loss_fn is not None` 检查 | 支持消融实验中将 gradient_loss_fn 设为 null | + +### 新增的文件 + +| 文件 | 用途 | +|------|------| +| `training/config/full_finetune.yaml` | 全参微调配置 (取消 Aggregator 冻结) | +| `training/config/ablation_no_depth.yaml` | 消融实验: 移除深度损失 | +| `training/config/ablation_no_camera.yaml` | 消融实验: 移除相机损失 | +| `training/config/ablation_no_grad_depth.yaml` | 消融实验: 移除深度梯度损失 | +| `training/launch_multi.py` | 灵活的启动器，支持 `--config` 参数选择配置 | +| `inference.py` | 样本外数据推理脚本 | +| `hw_monitor.py` | 硬件使用监控脚本 | +| `run_experiments.sh` | 一键运行所有实验的 shell 脚本 | +| `ASSIGNMENT_GUIDE.md` | 本文档 | + +### 运行全部任务的快速指南 + +```bash +# 1. 环境准备 +pip install -r requirements.txt + +# 2. 配置数据路径 (编辑 run_experiments.sh 或使用环境变量) +export CO3D_DIR="/your/path/to/co3d" +export CO3D_ANNO_DIR="/your/path/to/co3d_anno" +export PRETRAINED_CKPT="/your/path/to/model.pt" + +# 3. 运行基线实验 (部分微调) +torchrun --nproc_per_node=1 training/launch_multi.py --config default + +# 4. 运行消融实验 +torchrun --nproc_per_node=1 training/launch_multi.py --config ablation_no_grad_depth +torchrun --nproc_per_node=1 training/launch_multi.py --config ablation_no_depth +torchrun --nproc_per_node=1 training/launch_multi.py --config ablation_no_camera + +# 5. 运行全参微调 (需要 H100!) +torchrun --nproc_per_node=1 training/launch_multi.py --config full_finetune + +# 6. 推理测试 +python inference.py --image_dir ./test_images --output_dir ./vggt_output + +# 7. 查看 TensorBoard +tensorboard --logdir logs/tensorboard + +# 8. 查看硬件报告 +python hw_monitor.py --report hw_stats.json +``` + +--- + +## 文件结构总结 + +``` +vggt_training/ +├── vggt/ # VGGT 模型代码 +│ ├── models/ +│ │ ├── vggt.py # 主模型 (VGGT class) +│ │ └── aggregator.py # Aggregator (交错注意力) +│ ├── heads/ +│ │ ├── camera_head.py # 相机位姿预测头 +│ │ ├── dpt_head.py # 深度/点云 DPT 头 +│ │ └── track_head.py # 跟踪预测头 +│ └── layers/ # 基础层 (Attention, ViT, etc.) +├── training/ # 训练代码 +│ ├── config/ +│ │ ├── default.yaml # 默认配置 (部分微调) +│ │ ├── full_finetune.yaml # [NEW] 全参微调配置 +│ │ ├── ablation_no_depth.yaml # [NEW] 消融: 无深度损失 +│ │ ├── ablation_no_camera.yaml # [NEW] 消融: 无相机损失 +│ │ └── ablation_no_grad_depth.yaml # [NEW] 消融: 无梯度损失 +│ ├── data/ # 数据加载 +│ ├── train_utils/ # 训练工具 +│ ├── trainer.py # [MODIFIED] 主训练器 +│ ├── loss.py # [MODIFIED] 损失函数 +│ ├── launch.py # [MODIFIED] 原始启动器 +│ └── launch_multi.py # [NEW] 灵活启动器 +├── inference.py # [NEW] 推理脚本 +├── hw_monitor.py # [NEW] 硬件监控 +├── run_experiments.sh # [NEW] 实验运行脚本 +└── ASSIGNMENT_GUIDE.md # [NEW] 本文档 +``` diff --git a/hw_monitor.py b/hw_monitor.py new file mode 100644 index 000000000..64d543696 --- /dev/null +++ b/hw_monitor.py @@ -0,0 +1,264 @@ +""" +VGGT 训练/推理硬件监控脚本 + +功能: + 1. 记录 GPU 显存使用、利用率、温度等指标 + 2. 记录 CPU 内存使用 + 3. 支持后台运行，定时采样 + 4. 输出硬件使用报告 + +用法: + # 后台监控训练过程（每 2 秒采样一次，保存到文件） + python hw_monitor.py --interval 2 --output hw_stats.json + + # 在另一个终端运行训练，监控会自动记录 + + # 读取并打印监控报告 + python hw_monitor.py --report hw_stats.json +""" + +import argparse +import json +import os +import time +import threading +import subprocess +import sys +from datetime import datetime + + +def get_gpu_info(): + """获取 GPU 信息（通过 nvidia-smi）""" + try: + result = subprocess.run( + ["nvidia-smi", "--query-gpu=index,name,utilization.gpu,utilization.memory," + "memory.used,memory.total,memory.free,temperature.gpu,power.draw,power.limit", + "--format=csv,noheader,nounits"], + capture_output=True, text=True, timeout=5 + ) + gpus = [] + for line in result.stdout.strip().split('\n'): + if line.strip(): + parts = [p.strip() for p in line.split(',')] + if len(parts) >= 10: + gpus.append({ + "index": int(parts[0]), + "name": parts[1], + "gpu_util_pct": float(parts[2]), + "mem_util_pct": float(parts[3]), + "mem_used_mb": float(parts[4]), + "mem_total_mb": float(parts[5]), + "mem_free_mb": float(parts[6]), + "temp_c": float(parts[7]), + "power_w": float(parts[8]) if parts[8] != 'N/A' else None, + "power_limit_w": float(parts[9]) if parts[9] != 'N/A' else None, + }) + return gpus + except Exception as e: + return [{"error": str(e)}] + + +def get_cpu_memory(): + """获取 CPU 内存信息""" + try: + import psutil + mem = psutil.virtual_memory() + return { + "total_gb": mem.total / (1024**3), + "used_gb": mem.used / (1024**3), + "available_gb": mem.available / (1024**3), + "percent": mem.percent, + } + except ImportError: + return {"error": "psutil not installed. Run: pip install psutil"} + + +def get_pytorch_gpu_info(): + """获取 PyTorch 报告的 GPU 显存信息""" + try: + import torch + if torch.cuda.is_available(): + return { + "allocated_gb": torch.cuda.memory_allocated() / 1024**3, + "reserved_gb": torch.cuda.memory_reserved() / 1024**3, + "max_allocated_gb": torch.cuda.max_memory_allocated() / 1024**3, + "max_reserved_gb": torch.cuda.max_memory_reserved() / 1024**3, + } + except ImportError: + pass + return None + + +class HWMonitor: + """硬件监控器，在后台线程中定时采样""" + + def __init__(self, interval=2.0): + self.interval = interval + self.samples = [] + self.running = False + self.thread = None + + def start(self): + """启动后台监控""" + self.running = True + self.thread = threading.Thread(target=self._monitor_loop, daemon=True) + self.thread.start() + print(f"[HW Monitor] Started (interval={self.interval}s)") + + def stop(self): + """停止监控""" + self.running = False + if self.thread: + self.thread.join(timeout=5) + print(f"[HW Monitor] Stopped ({len(self.samples)} samples)") + + def _monitor_loop(self): + """监控循环""" + while self.running: + sample = { + "timestamp": datetime.now().isoformat(), + "gpus": get_gpu_info(), + "cpu_mem": get_cpu_memory(), + "pytorch_gpu": get_pytorch_gpu_info(), + } + self.samples.append(sample) + time.sleep(self.interval) + + def save(self, output_path): + """保存监控数据到文件""" + with open(output_path, 'w') as f: + json.dump(self.samples, f, indent=2) + print(f"[HW Monitor] Data saved to {output_path}") + + def generate_report(self): + """生成硬件使用报告""" + if not self.samples: + return "No samples recorded." + + # 提取 GPU 指标序列 + gpu_metrics = {} + for sample in self.samples: + for gpu in sample.get("gpus", []): + idx = gpu.get("index", 0) + if idx not in gpu_metrics: + gpu_metrics[idx] = { + "name": gpu.get("name", "Unknown"), + "mem_used": [], + "mem_total": [], + "gpu_util": [], + "temp": [], + "power": [], + "timestamps": [], + } + gpu_metrics[idx]["mem_used"].append(gpu.get("mem_used_mb", 0)) + gpu_metrics[idx]["mem_total"].append(gpu.get("mem_total_mb", 0)) + gpu_metrics[idx]["gpu_util"].append(gpu.get("gpu_util_pct", 0)) + gpu_metrics[idx]["temp"].append(gpu.get("temp_c", 0)) + gpu_metrics[idx]["power"].append(gpu.get("power_w", 0) or 0) + gpu_metrics[idx]["timestamps"].append(sample["timestamp"]) + + # 生成报告 + report = [] + report.append("=" * 70) + report.append("VGGT 硬件使用报告") + report.append("=" * 70) + report.append(f"采样时间: {self.samples[0]['timestamp']} -> {self.samples[-1]['timestamp']}") + report.append(f"采样次数: {len(self.samples)}") + report.append(f"采样间隔: {self.interval}s") + report.append("") + + for idx, metrics in sorted(gpu_metrics.items()): + report.append(f"--- GPU {idx}: {metrics['name']} ---") + mem_used = metrics["mem_used"] + mem_total = metrics["mem_total"] + report.append(f" 显存总量: {mem_total[0]:.0f} MB") + report.append(f" 显存峰值: {max(mem_used):.0f} MB ({max(mem_used)/mem_total[0]*100:.1f}%)") + report.append(f" 显存均值: {sum(mem_used)/len(mem_used):.0f} MB") + report.append(f" GPU 利用率峰值: {max(metrics['gpu_util']):.1f}%") + report.append(f" GPU 利用率均值: {sum(metrics['gpu_util'])/len(metrics['gpu_util']):.1f}%") + report.append(f" 温度峰值: {max(metrics['temp']):.1f}°C") + report.append(f" 功耗峰值: {max(metrics['power']):.1f}W") + report.append("") + + # PyTorch 显存 + pytorch_mem = [s.get("pytorch_gpu", {}) for s in self.samples if s.get("pytorch_gpu")] + if pytorch_mem: + max_alloc = max(m.get("max_allocated_gb", 0) or 0 for m in pytorch_mem if m) + max_resv = max(m.get("max_reserved_gb", 0) or 0 for m in pytorch_mem if m) + report.append("--- PyTorch 显存统计 ---") + report.append(f" 峰值分配 (allocated): {max_alloc:.2f} GB") + report.append(f" 峰值保留 (reserved): {max_resv:.2f} GB") + report.append("") + + # CPU 内存 + cpu_samples = [s.get("cpu_mem", {}) for s in self.samples if s.get("cpu_mem") and "error" not in s["cpu_mem"]] + if cpu_samples: + cpu_max_pct = max(s.get("percent", 0) for s in cpu_samples) + cpu_max_used = max(s.get("used_gb", 0) for s in cpu_samples) + report.append("--- CPU 内存 ---") + report.append(f" 峰值使用: {cpu_max_used:.2f} GB ({cpu_max_pct:.1f}%)") + report.append("") + + return "\n".join(report) + + +def print_report_from_file(filepath): + """从保存的文件中读取并打印报告""" + with open(filepath, 'r') as f: + samples = json.load(f) + + monitor = HWMonitor() + monitor.samples = samples + # 从采样间隔推断 + if len(samples) >= 2: + t1 = datetime.fromisoformat(samples[0]["timestamp"]) + t2 = datetime.fromisoformat(samples[1]["timestamp"]) + monitor.interval = (t2 - t1).total_seconds() + print(monitor.generate_report()) + + +def main(): + parser = argparse.ArgumentParser(description="VGGT 硬件监控工具") + parser.add_argument("--interval", type=float, default=2.0, + help="采样间隔（秒）") + parser.add_argument("--output", type=str, default="hw_stats.json", + help="输出文件路径") + parser.add_argument("--duration", type=float, default=None, + help="监控时长（秒），默认持续运行直到 Ctrl+C") + parser.add_argument("--report", type=str, default=None, + help="从已有的 JSON 文件生成报告") + args = parser.parse_args() + + # 报告模式 + if args.report: + print_report_from_file(args.report) + return + + # 监控模式 + monitor = HWMonitor(interval=args.interval) + monitor.start() + + print(f"\n{'='*60}") + print("硬件监控运行中...") + print(f"输出文件: {args.output}") + if args.duration: + print(f"将运行 {args.duration} 秒") + print("按 Ctrl+C 停止") + print(f"{'='*60}\n") + + try: + if args.duration: + time.sleep(args.duration) + else: + while True: + time.sleep(1) + except KeyboardInterrupt: + print("\n正在停止监控...") + + monitor.stop() + monitor.save(args.output) + print(monitor.generate_report()) + + +if __name__ == "__main__": + main() diff --git a/inference.py b/inference.py new file mode 100644 index 000000000..1cf9da20f --- /dev/null +++ b/inference.py @@ -0,0 +1,302 @@ +""" +VGGT 推理脚本 - 用于对样本外（out-of-sample）数据进行推理 + +功能: + 1. 加载预训练/微调后的 VGGT 模型 + 2. 对任意文件夹中的图像进行 3D 重建 + 3. 输出相机位姿、深度图、3D点云等结果 + +用法: + # 使用预训练模型（从 HuggingFace 下载） + python inference.py --image_dir /path/to/images --output_dir /path/to/output + + # 使用微调后的 checkpoint + python inference.py --image_dir /path/to/images --output_dir /path/to/output --checkpoint logs/exp001/ckpts/checkpoint.pt + + # 使用 BA (Bundle Adjustment) 优化 + python inference.py --image_dir /path/to/images --output_dir /path/to/output --use_ba +""" + +import os +import sys +import argparse +import glob +import time +import json +import numpy as np +import torch +import torch.nn.functional as F + +# Add project root to Python path +_project_root = os.path.dirname(os.path.abspath(__file__)) +if _project_root not in sys.path: + sys.path.insert(0, _project_root) + +from vggt.models.vggt import VGGT +from vggt.utils.load_fn import load_and_preprocess_images_square +from vggt.utils.pose_enc import pose_encoding_to_extri_intri +from vggt.utils.geometry import unproject_depth_map_to_point_map +from vggt.utils.helper import create_pixel_coordinate_grid, randomly_limit_trues + + +def parse_args(): + parser = argparse.ArgumentParser(description="VGGT Inference Script") + parser.add_argument("--image_dir", type=str, required=True, + help="输入图像目录") + parser.add_argument("--output_dir", type=str, default="./vggt_output", + help="输出目录") + parser.add_argument("--checkpoint", type=str, default=None, + help="微调后的 checkpoint 路径（可选，默认使用 HuggingFace 预训练模型）") + parser.add_argument("--resolution", type=int, default=518, + help="VGGT 输入分辨率（默认 518）") + parser.add_argument("--conf_thres", type=float, default=5.0, + help="深度置信度阈值") + parser.add_argument("--seed", type=int, default=42, + help="随机种子") + return parser.parse_args() + + +def setup_device_and_dtype(): + """设置设备和数据类型""" + if torch.cuda.is_available(): + device = "cuda" + capability = torch.cuda.get_device_capability() + dtype = torch.bfloat16 if capability[0] >= 8 else torch.float16 + print(f"GPU: {torch.cuda.get_device_name()}") + print(f"Compute Capability: {capability}") + else: + device = "cpu" + dtype = torch.float32 + print("WARNING: No GPU detected, using CPU (will be slow)") + print(f"Device: {device}, Dtype: {dtype}") + return device, dtype + + +def load_model(device, checkpoint_path=None): + """加载 VGGT 模型""" + print("\n" + "=" * 60) + print("加载 VGGT 模型...") + + model = VGGT() + + if checkpoint_path is not None and os.path.exists(checkpoint_path): + # 加载微调后的 checkpoint + print(f"从 {checkpoint_path} 加载微调权重...") + checkpoint = torch.load(checkpoint_path, map_location="cpu") + if "model" in checkpoint: + state_dict = checkpoint["model"] + else: + state_dict = checkpoint + model.load_state_dict(state_dict, strict=False) + print("微调权重加载完成") + else: + # 从 HuggingFace 下载预训练模型 + print("从 HuggingFace 下载预训练模型...") + _URL = "https://huggingface.co/facebook/VGGT-1B/resolve/main/model.pt" + model.load_state_dict(torch.hub.load_state_dict_from_url(_URL)) + print("预训练模型加载完成") + + model.eval() + model = model.to(device) + + # 统计参数量 + total_params = sum(p.numel() for p in model.parameters()) + trainable_params = sum(p.numel() for p in model.parameters() if p.requires_grad) + print(f"总参数量: {total_params:,}") + print(f"可训练参数量: {trainable_params:,}") + + return model + + +@torch.no_grad() +def run_inference(model, images, device, dtype, resolution=518): + """ + 运行 VGGT 推理 + + Args: + model: VGGT 模型 + images: (B, 3, H, W) 输入图像, 范围 [0, 1] + device: 设备 + dtype: 数据类型 + resolution: VGGT 输入分辨率 + + Returns: + extrinsic: (B, 3, 4) 相机外参 + intrinsic: (B, 3, 3) 相机内参 + depth_map: (B, H, W, 1) 深度图 + depth_conf: (B, H, W) 深度置信度 + points_3d: (B, H, W, 3) 3D 点云 + """ + B = images.shape[0] + + # Resize to VGGT resolution + images_resized = F.interpolate( + images, size=(resolution, resolution), mode="bilinear", align_corners=False + ) + images_resized = images_resized[None] # add batch dim: (1, B, 3, H, W) + + with torch.cuda.amp.autocast(dtype=dtype): + # 通过 Aggregator 提取特征 + aggregated_tokens_list, ps_idx = model.aggregator(images_resized) + + # 预测相机位姿 + if model.camera_head is not None: + pose_enc = model.camera_head(aggregated_tokens_list)[-1] + extrinsic, intrinsic = pose_encoding_to_extri_intri(pose_enc, images_resized.shape[-2:]) + else: + extrinsic = None + intrinsic = None + + # 预测深度图 + if model.depth_head is not None: + depth_map, depth_conf = model.depth_head(aggregated_tokens_list, images_resized, ps_idx) + else: + depth_map = None + depth_conf = None + + # 后处理 + if extrinsic is not None: + extrinsic = extrinsic.squeeze(0).cpu().numpy() + intrinsic = intrinsic.squeeze(0).cpu().numpy() + + if depth_map is not None: + depth_map = depth_map.squeeze(0).cpu().numpy() + depth_conf = depth_conf.squeeze(0).cpu().numpy() + # Unproject depth to 3D points + if extrinsic is not None: + points_3d = unproject_depth_map_to_point_map(depth_map, extrinsic, intrinsic) + else: + points_3d = None + else: + points_3d = None + + return extrinsic, intrinsic, depth_map, depth_conf, points_3d + + +def save_results(output_dir, image_names, extrinsic, intrinsic, depth_map, depth_conf, points_3d, conf_thres): + """保存推理结果""" + os.makedirs(output_dir, exist_ok=True) + + # 保存相机位姿 + if extrinsic is not None and intrinsic is not None: + pose_data = { + "extrinsics": extrinsic.tolist(), + "intrinsics": intrinsic.tolist(), + "image_names": image_names, + } + with open(os.path.join(output_dir, "cameras.json"), "w") as f: + json.dump(pose_data, f, indent=2) + print(f"相机参数已保存至 {output_dir}/cameras.json") + + # 保存深度图统计信息 + if depth_map is not None: + depth_stats = { + "mean_depth": float(depth_map[depth_map > 0].mean()), + "median_depth": float(np.median(depth_map[depth_map > 0])), + "min_depth": float(depth_map[depth_map > 0].min()), + "max_depth": float(depth_map.max()), + "mean_confidence": float(depth_conf.mean()), + } + with open(os.path.join(output_dir, "depth_stats.json"), "w") as f: + json.dump(depth_stats, f, indent=2) + print(f"深度统计已保存至 {output_dir}/depth_stats.json") + + # 保存深度图为 .npy 文件 + np.save(os.path.join(output_dir, "depth_maps.npy"), depth_map) + np.save(os.path.join(output_dir, "depth_confs.npy"), depth_conf) + print(f"深度图已保存至 {output_dir}/depth_maps.npy") + + # 保存 3D 点云（仅高置信度部分） + if points_3d is not None and depth_conf is not None: + conf_mask = depth_conf >= conf_thres + valid_points = points_3d[conf_mask] + + # 随机采样限制点数 + max_points = 100000 + if len(valid_points) > max_points: + indices = np.random.choice(len(valid_points), max_points, replace=False) + valid_points = valid_points[indices] + + np.save(os.path.join(output_dir, "points_3d.npy"), valid_points) + print(f"3D 点云已保存至 {output_dir}/points_3d.npy ({len(valid_points)} 个点)") + + print(f"\n所有结果已保存至 {output_dir}/") + print("=" * 60) + + +def main(): + args = parse_args() + + # 设置随机种子 + np.random.seed(args.seed) + torch.manual_seed(args.seed) + + # 设置设备 + device, dtype = setup_device_and_dtype() + + # 获取图像列表 + image_extensions = ["*.jpg", "*.jpeg", "*.png", "*.JPG", "*.JPEG", "*.PNG"] + image_paths = [] + for ext in image_extensions: + image_paths.extend(glob.glob(os.path.join(args.image_dir, ext))) + image_paths = sorted(image_paths) + + if len(image_paths) == 0: + raise ValueError(f"在 {args.image_dir} 中未找到图像文件") + + print(f"\n找到 {len(image_paths)} 张图像") + image_names = [os.path.basename(p) for p in image_paths] + + # 加载并预处理图像 + print("加载并预处理图像...") + images, original_coords = load_and_preprocess_images_square( + image_paths, img_load_resolution=1024 + ) + images = images.to(device) + print(f"图像张量形状: {images.shape}") + + # 加载模型 + model = load_model(device, args.checkpoint) + + # 运行推理 + print("\n运行 VGGT 推理...") + start_time = time.time() + + extrinsic, intrinsic, depth_map, depth_conf, points_3d = run_inference( + model, images, device, dtype, args.resolution + ) + + elapsed = time.time() - start_time + print(f"推理完成，耗时 {elapsed:.2f} 秒") + + # 打印结果摘要 + print("\n" + "=" * 60) + print("推理结果摘要:") + if extrinsic is not None: + print(f" 相机外参形状: {extrinsic.shape}") + print(f" 相机内参形状: {intrinsic.shape}") + # 打印第一帧的相机位置 + R = extrinsic[0, :3, :3] + t = extrinsic[0, :3, 3] + cam_position = -R.T @ t + print(f" 第一帧相机位置 (世界坐标): {cam_position}") + if depth_map is not None: + print(f" 深度图形状: {depth_map.shape}") + print(f" 深度置信度形状: {depth_conf.shape}") + valid_depth_pixels = (depth_map > 0).sum() + print(f" 有效深度像素数: {valid_depth_pixels:,}") + + # 保存结果 + print("\n保存结果...") + save_results(args.output_dir, image_names, extrinsic, intrinsic, + depth_map, depth_conf, points_3d, args.conf_thres) + + # 记录显存使用 + if torch.cuda.is_available(): + print(f"\n显存使用:") + print(f" Allocated: {torch.cuda.max_memory_allocated() / 1024**3:.2f} GB") + print(f" Reserved: {torch.cuda.max_memory_reserved() / 1024**3:.2f} GB") + + +if __name__ == "__main__": + main() diff --git a/run_experiments.sh b/run_experiments.sh new file mode 100644 index 000000000..1f5ad7f73 --- /dev/null +++ b/run_experiments.sh @@ -0,0 +1,188 @@ +#!/bin/bash +# ============================================================================= +# VGGT 实验运行脚本 +# 包含: 全参微调 + 消融实验 + 推理 + 硬件监控 +# +# 用法: +# chmod +x run_experiments.sh +# ./run_experiments.sh +# +# 环境要求: +# - PyTorch >= 2.0 +# - CUDA >= 11.8 +# - GPU: H100 (80GB) for full fine-tuning; 4090 (24GB) for partial fine-tuning +# - Co3D 数据集（需提前下载并配置路径） +# ============================================================================= + +set -e # 遇到错误立即退出 + +# ============================================================================= +# 配置区 - 请根据实际环境修改 +# ============================================================================= + +# Co3D 数据路径（修改为你实际的路径） +CO3D_DIR="${CO3D_DIR:-/fsx-repligen/jianyuan/transfer_buffer/small_set/co3d/}" +CO3D_ANNO_DIR="${CO3D_ANNO_DIR:-/fsx-repligen/jianyuan/transfer_buffer/small_set/co3d_anno}" + +# 预训练 checkpoint 路径 +PRETRAINED_CKPT="${PRETRAINED_CKPT:-/fsx-repligen/jianyuan/transfer_buffer/ckpts/model.pt}" + +# GPU 数量（1 为单卡，>1 为多卡 DDP） +NUM_GPUS="${NUM_GPUS:-1}" + +# 推理测试图像目录 +INFERENCE_IMAGE_DIR="${INFERENCE_IMAGE_DIR:-./test_images}" + +# ============================================================================= +# 工具函数 +# ============================================================================= + +run_training() { + local config_name=$1 + local description=$2 + + echo "" + echo "#######################################################################" + echo "# 实验: ${description}" + echo "# 配置: ${config_name}.yaml" + echo "#######################################################################" + + # 启动硬件监控 + local hw_log="logs/hw_${config_name}.json" + python hw_monitor.py --interval 5 --output "${hw_log}" & + local monitor_pid=$! + + # 运行训练 + torchrun --nproc_per_node=${NUM_GPUS} training/launch_multi.py \ + --config "${config_name}" \ + --override \ + "data.train.dataset.dataset_configs.0.CO3D_DIR=${CO3D_DIR}" \ + "data.train.dataset.dataset_configs.0.CO3D_ANNOTATION_DIR=${CO3D_ANNO_DIR}" \ + "data.val.dataset.dataset_configs.0.CO3D_DIR=${CO3D_DIR}" \ + "data.val.dataset.dataset_configs.0.CO3D_ANNOTATION_DIR=${CO3D_ANNO_DIR}" \ + "checkpoint.resume_checkpoint_path=${PRETRAINED_CKPT}" + + # 停止硬件监控 + kill ${monitor_pid} 2>/dev/null || true + wait ${monitor_pid} 2>/dev/null || true + + # 生成硬件报告 + echo "" + echo "--- 硬件使用报告: ${config_name} ---" + python hw_monitor.py --report "${hw_log}" + + echo "实验 ${config_name} 完成！" +} + +run_inference() { + local checkpoint_path=$1 + local description=$2 + + echo "" + echo "#######################################################################" + echo "# 推理: ${description}" + echo "#######################################################################" + + local output_dir="vggt_output/${description// /_}" + + python inference.py \ + --image_dir "${INFERENCE_IMAGE_DIR}" \ + --output_dir "${output_dir}" \ + --checkpoint "${checkpoint_path}" \ + --conf_thres 5.0 +} + +# ============================================================================= +# 主流程 +# ============================================================================= + +echo "" +echo "╔══════════════════════════════════════════════════════════════════════╗" +echo "║ VGGT 实验运行脚本 ║" +echo "╠══════════════════════════════════════════════════════════════════════╣" +echo "║ GPU 数量: ${NUM_GPUS}" +echo "║ Co3D 数据目录: ${CO3D_DIR}" +echo "║ Co3D 标注目录: ${CO3D_ANNO_DIR}" +echo "║ 预训练 Checkpoint: ${PRETRAINED_CKPT}" +echo "╚══════════════════════════════════════════════════════════════════════╝" +echo "" + +# 选择要运行的实验（根据需求取消注释） +SELECTED_EXPERIMENT="${1:-all}" + +case ${SELECTED_EXPERIMENT} in + baseline) + # 实验 0: 基线（默认配置：部分微调，冻结 Aggregator） + run_training "default" "Baseline (Partial Fine-tuning)" + ;; + + full) + # 实验 1: 全参微调（需要 H100 80GB) + echo "WARNING: 全参微调需要 H100 80GB 或以上 GPU!" + echo "4090 (24GB) 无法运行此实验" + echo "继续? (y/n)" + read -r confirm + if [ "${confirm}" != "y" ]; then + echo "跳过全参微调" + else + run_training "full_finetune" "Full-Parameter Fine-tuning" + fi + ;; + + ablation) + # 实验 2: 消融实验组 + echo "运行消融实验组..." + + # 消融 2a: 移除深度损失 + run_training "ablation_no_depth" "Ablation: No Depth Loss" + + # 消融 2b: 移除相机损失 + run_training "ablation_no_camera" "Ablation: No Camera Loss" + + # 消融 2c: 移除深度梯度损失 + run_training "ablation_no_grad_depth" "Ablation: No Depth Gradient Loss" + ;; + + inference) + # 推理测试 + run_inference "" "Pretrained Model Inference" + + # 如果有微调后的模型，也进行推理 + if [ -f "logs/exp001/ckpts/checkpoint.pt" ]; then + run_inference "logs/exp001/ckpts/checkpoint.pt" "Fine-tuned Model Inference" + fi + ;; + + all) + echo "运行全部实验..." + # 1. 基线 + run_training "default" "Baseline" + + # 2. 消融实验 + run_training "ablation_no_grad_depth" "Ablation: No Depth Gradient Loss" + run_training "ablation_no_depth" "Ablation: No Depth Loss" + run_training "ablation_no_camera" "Ablation: No Camera Loss" + + # 3. 推理测试 + run_inference "" "Pretrained Model" + ;; + + *) + echo "用法: $0 [baseline|full|ablation|inference|all]" + echo "" + echo " baseline - 基线训练（部分微调）" + echo " full - 全参微调（需要 H100）" + echo " ablation - 消融实验组（3 个实验）" + echo " inference - 推理测试" + echo " all - 全部实验" + exit 1 + ;; +esac + +echo "" +echo "╔══════════════════════════════════════════════════════════════════════╗" +echo "║ 所有实验完成！ ║" +echo "║ TensorBoard: tensorboard --logdir logs/tensorboard ║" +echo "║ 日志目录: logs/ ║" +echo "║ Checkpoint 目录: logs/*/ckpts/ ║" +echo "╚══════════════════════════════════════════════════════════════════════╝" diff --git a/training/config/ablation_no_camera.yaml b/training/config/ablation_no_camera.yaml new file mode 100644 index 000000000..c72d72ef1 --- /dev/null +++ b/training/config/ablation_no_camera.yaml @@ -0,0 +1,157 @@ +# ============================================================================= +# Ablation 2: Depth Loss Only (Camera Loss removed) +# Tests whether camera pose loss can be removed during fine-tuning +# ============================================================================= + +defaults: + - default_dataset.yaml + +exp_name: ablation_no_camera +img_size: 518 +num_workers: 0 +seed_value: 42 +accum_steps: 3 +patch_size: 14 + +limit_train_batches: 800 +limit_val_batches: 400 + +data: + train: + _target_: data.dynamic_dataloader.DynamicTorchDataset + num_workers: ${num_workers} + common_config: + img_size: ${img_size} + patch_size: ${patch_size} + debug: True + repeat_batch: True + dataset: + _target_: data.composed_dataset.ComposedDataset + dataset_configs: + - _target_: data.datasets.co3d.Co3dDataset + split: train + CO3D_DIR: /fsx-repligen/jianyuan/transfer_buffer/small_set/co3d/ + CO3D_ANNOTATION_DIR: /fsx-repligen/jianyuan/transfer_buffer/small_set/co3d_anno + val: + _target_: data.dynamic_dataloader.DynamicTorchDataset + num_workers: ${num_workers} + common_config: + img_size: ${img_size} + patch_size: ${patch_size} + debug: True + dataset: + _target_: data.composed_dataset.ComposedDataset + dataset_configs: + - _target_: data.datasets.co3d.Co3dDataset + split: test + CO3D_DIR: /fsx-repligen/jianyuan/transfer_buffer/small_set/co3d/ + CO3D_ANNOTATION_DIR: /fsx-repligen/jianyuan/transfer_buffer/small_set/co3d_anno + + +logging: + log_dir: logs + log_visuals: False + log_freq: 1 + log_level_primary: DEBUG + log_level_secondary: WARNING + all_ranks: False + tensorboard_writer: + _target_: train_utils.tb_writer.TensorBoardLogger + path: ${logging.log_dir}/tensorboard + scalar_keys_to_log: + train: + keys_to_log: + - loss_objective + - loss_conf_depth + - loss_reg_depth + - loss_grad_depth + val: + keys_to_log: + - loss_objective + - loss_conf_depth + - loss_reg_depth + - loss_grad_depth + + + +checkpoint: + save_dir: logs/${exp_name}/ckpts + save_freq: 5 + resume_checkpoint_path: /fsx-repligen/jianyuan/transfer_buffer/ckpts/model.pt + strict: False + + +loss: + _target_: loss.MultitaskLoss + # KEY CHANGE: Remove camera loss entirely + camera: null + depth: + weight: 1.0 + gradient_loss_fn: "grad" + valid_range: 0.98 + point: null + track: null + + +optim: + param_group_modifiers: False + optimizer: + _target_: torch.optim.AdamW + lr: 1e-4 + weight_decay: 0.05 + frozen_module_names: + - "*aggregator*" + amp: + enabled: True + amp_dtype: bfloat16 + gradient_clip: + _target_: train_utils.gradient_clip.GradientClipper + configs: + - module_name: ["aggregator"] + max_norm: 1.0 + norm_type: 2 + - module_name: ["depth"] + max_norm: 1.0 + norm_type: 2 + options: + lr: + - scheduler: + _target_: fvcore.common.param_scheduler.CompositeParamScheduler + schedulers: + - _target_: fvcore.common.param_scheduler.LinearParamScheduler + start_value: 1e-8 + end_value: 1e-4 + - _target_: fvcore.common.param_scheduler.CosineParamScheduler + start_value: 1e-4 + end_value: 1e-8 + lengths: [0.05, 0.95] + interval_scaling: ['rescaled', 'rescaled'] + weight_decay: + - scheduler: + _target_: fvcore.common.param_scheduler.ConstantParamScheduler + value: 0.05 + + +max_epochs: 100 + +model: + _target_: vggt.models.vggt.VGGT + enable_camera: False # Disable camera prediction head + enable_depth: True + enable_point: False + enable_track: False + + +distributed: + backend: nccl + comms_dtype: None + find_unused_parameters: False + timeout_mins: 30 + gradient_as_bucket_view: True + bucket_cap_mb: 25 + broadcast_buffers: True + +cuda: + cudnn_deterministic: False + cudnn_benchmark: False + allow_tf32: True diff --git a/training/config/ablation_no_depth.yaml b/training/config/ablation_no_depth.yaml new file mode 100644 index 000000000..f6890b356 --- /dev/null +++ b/training/config/ablation_no_depth.yaml @@ -0,0 +1,158 @@ +# ============================================================================= +# Ablation 1: Camera Loss Only (Depth Loss removed) +# Tests whether depth loss can be removed during fine-tuning +# ============================================================================= + +defaults: + - default_dataset.yaml + +exp_name: ablation_no_depth +img_size: 518 +num_workers: 0 +seed_value: 42 +accum_steps: 3 +patch_size: 14 + +limit_train_batches: 800 +limit_val_batches: 400 + +data: + train: + _target_: data.dynamic_dataloader.DynamicTorchDataset + num_workers: ${num_workers} + common_config: + img_size: ${img_size} + patch_size: ${patch_size} + debug: True + repeat_batch: True + dataset: + _target_: data.composed_dataset.ComposedDataset + dataset_configs: + - _target_: data.datasets.co3d.Co3dDataset + split: train + CO3D_DIR: /fsx-repligen/jianyuan/transfer_buffer/small_set/co3d/ + CO3D_ANNOTATION_DIR: /fsx-repligen/jianyuan/transfer_buffer/small_set/co3d_anno + val: + _target_: data.dynamic_dataloader.DynamicTorchDataset + num_workers: ${num_workers} + common_config: + img_size: ${img_size} + patch_size: ${patch_size} + debug: True + dataset: + _target_: data.composed_dataset.ComposedDataset + dataset_configs: + - _target_: data.datasets.co3d.Co3dDataset + split: test + CO3D_DIR: /fsx-repligen/jianyuan/transfer_buffer/small_set/co3d/ + CO3D_ANNOTATION_DIR: /fsx-repligen/jianyuan/transfer_buffer/small_set/co3d_anno + + +logging: + log_dir: logs + log_visuals: False + log_freq: 1 + log_level_primary: DEBUG + log_level_secondary: WARNING + all_ranks: False + tensorboard_writer: + _target_: train_utils.tb_writer.TensorBoardLogger + path: ${logging.log_dir}/tensorboard + scalar_keys_to_log: + train: + keys_to_log: + - loss_objective + - loss_camera + - loss_T + - loss_R + - loss_FL + val: + keys_to_log: + - loss_objective + - loss_camera + - loss_T + - loss_R + - loss_FL + + + +checkpoint: + save_dir: logs/${exp_name}/ckpts + save_freq: 5 + resume_checkpoint_path: /fsx-repligen/jianyuan/transfer_buffer/ckpts/model.pt + strict: False + + +loss: + _target_: loss.MultitaskLoss + camera: + weight: 5.0 + loss_type: "l1" + # KEY CHANGE: Remove depth loss entirely + depth: null + point: null + track: null + + +optim: + param_group_modifiers: False + optimizer: + _target_: torch.optim.AdamW + lr: 1e-4 + weight_decay: 0.05 + frozen_module_names: + - "*aggregator*" + amp: + enabled: True + amp_dtype: bfloat16 + gradient_clip: + _target_: train_utils.gradient_clip.GradientClipper + configs: + - module_name: ["aggregator"] + max_norm: 1.0 + norm_type: 2 + - module_name: ["camera"] + max_norm: 1.0 + norm_type: 2 + options: + lr: + - scheduler: + _target_: fvcore.common.param_scheduler.CompositeParamScheduler + schedulers: + - _target_: fvcore.common.param_scheduler.LinearParamScheduler + start_value: 1e-8 + end_value: 1e-4 + - _target_: fvcore.common.param_scheduler.CosineParamScheduler + start_value: 1e-4 + end_value: 1e-8 + lengths: [0.05, 0.95] + interval_scaling: ['rescaled', 'rescaled'] + weight_decay: + - scheduler: + _target_: fvcore.common.param_scheduler.ConstantParamScheduler + value: 0.05 + + +max_epochs: 100 + +model: + _target_: vggt.models.vggt.VGGT + enable_camera: True + enable_depth: False # Disable depth prediction head + enable_point: False + enable_track: False + + +distributed: + backend: nccl + comms_dtype: None + find_unused_parameters: False + timeout_mins: 30 + gradient_as_bucket_view: True + bucket_cap_mb: 25 + broadcast_buffers: True + +cuda: + cudnn_deterministic: False + cudnn_benchmark: False + allow_tf32: True diff --git a/training/config/ablation_no_grad_depth.yaml b/training/config/ablation_no_grad_depth.yaml new file mode 100644 index 000000000..282a63d58 --- /dev/null +++ b/training/config/ablation_no_grad_depth.yaml @@ -0,0 +1,170 @@ +# ============================================================================= +# Ablation 3: Remove Depth Gradient Loss (Camera + Depth without gradient) +# Tests whether the depth gradient smoothness loss can be removed +# ============================================================================= + +defaults: + - default_dataset.yaml + +exp_name: ablation_no_grad_depth +img_size: 518 +num_workers: 0 +seed_value: 42 +accum_steps: 3 +patch_size: 14 + +limit_train_batches: 800 +limit_val_batches: 400 + +data: + train: + _target_: data.dynamic_dataloader.DynamicTorchDataset + num_workers: ${num_workers} + common_config: + img_size: ${img_size} + patch_size: ${patch_size} + debug: True + repeat_batch: True + dataset: + _target_: data.composed_dataset.ComposedDataset + dataset_configs: + - _target_: data.datasets.co3d.Co3dDataset + split: train + CO3D_DIR: /fsx-repligen/jianyuan/transfer_buffer/small_set/co3d/ + CO3D_ANNOTATION_DIR: /fsx-repligen/jianyuan/transfer_buffer/small_set/co3d_anno + val: + _target_: data.dynamic_dataloader.DynamicTorchDataset + num_workers: ${num_workers} + common_config: + img_size: ${img_size} + patch_size: ${patch_size} + debug: True + dataset: + _target_: data.composed_dataset.ComposedDataset + dataset_configs: + - _target_: data.datasets.co3d.Co3dDataset + split: test + CO3D_DIR: /fsx-repligen/jianyuan/transfer_buffer/small_set/co3d/ + CO3D_ANNOTATION_DIR: /fsx-repligen/jianyuan/transfer_buffer/small_set/co3d_anno + + +logging: + log_dir: logs + log_visuals: False + log_freq: 1 + log_level_primary: DEBUG + log_level_secondary: WARNING + all_ranks: False + tensorboard_writer: + _target_: train_utils.tb_writer.TensorBoardLogger + path: ${logging.log_dir}/tensorboard + scalar_keys_to_log: + train: + keys_to_log: + - loss_objective + - loss_camera + - loss_T + - loss_R + - loss_FL + - loss_conf_depth + - loss_reg_depth + - loss_grad_depth + val: + keys_to_log: + - loss_objective + - loss_camera + - loss_T + - loss_R + - loss_FL + - loss_conf_depth + - loss_reg_depth + - loss_grad_depth + + + +checkpoint: + save_dir: logs/${exp_name}/ckpts + save_freq: 5 + resume_checkpoint_path: /fsx-repligen/jianyuan/transfer_buffer/ckpts/model.pt + strict: False + + +loss: + _target_: loss.MultitaskLoss + camera: + weight: 5.0 + loss_type: "l1" + depth: + weight: 1.0 + # KEY CHANGE: Remove gradient loss (set to null) + gradient_loss_fn: null + valid_range: 0.98 + point: null + track: null + + +optim: + param_group_modifiers: False + optimizer: + _target_: torch.optim.AdamW + lr: 1e-4 + weight_decay: 0.05 + frozen_module_names: + - "*aggregator*" + amp: + enabled: True + amp_dtype: bfloat16 + gradient_clip: + _target_: train_utils.gradient_clip.GradientClipper + configs: + - module_name: ["aggregator"] + max_norm: 1.0 + norm_type: 2 + - module_name: ["depth"] + max_norm: 1.0 + norm_type: 2 + - module_name: ["camera"] + max_norm: 1.0 + norm_type: 2 + options: + lr: + - scheduler: + _target_: fvcore.common.param_scheduler.CompositeParamScheduler + schedulers: + - _target_: fvcore.common.param_scheduler.LinearParamScheduler + start_value: 1e-8 + end_value: 1e-4 + - _target_: fvcore.common.param_scheduler.CosineParamScheduler + start_value: 1e-4 + end_value: 1e-8 + lengths: [0.05, 0.95] + interval_scaling: ['rescaled', 'rescaled'] + weight_decay: + - scheduler: + _target_: fvcore.common.param_scheduler.ConstantParamScheduler + value: 0.05 + + +max_epochs: 100 + +model: + _target_: vggt.models.vggt.VGGT + enable_camera: True + enable_depth: True + enable_point: False + enable_track: False + + +distributed: + backend: nccl + comms_dtype: None + find_unused_parameters: False + timeout_mins: 30 + gradient_as_bucket_view: True + bucket_cap_mb: 25 + broadcast_buffers: True + +cuda: + cudnn_deterministic: False + cudnn_benchmark: False + allow_tf32: True diff --git a/training/config/full_finetune.yaml b/training/config/full_finetune.yaml new file mode 100644 index 000000000..1567721b3 --- /dev/null +++ b/training/config/full_finetune.yaml @@ -0,0 +1,183 @@ +# ============================================================================= +# Full-Parameter Fine-tuning Configuration +# +# Key differences from default.yaml: +# 1. Removed frozen_module_names — all parameters are trainable +# 2. Increased accum_steps to compensate for smaller per-GPU batch +# 3. Requires H100 (80GB) or equivalent GPU +# +# Estimated GPU memory: ~60-70GB (H100 80GB) +# RTX 4090 (24GB) cannot support full-parameter fine-tuning +# ============================================================================= + +defaults: + - default_dataset.yaml + +exp_name: full_finetune +img_size: 518 +num_workers: 0 +seed_value: 42 +accum_steps: 4 # Increased gradient accumulation for full fine-tuning +patch_size: 14 + +limit_train_batches: 800 +limit_val_batches: 400 + +data: + train: + _target_: data.dynamic_dataloader.DynamicTorchDataset + num_workers: ${num_workers} + common_config: + img_size: ${img_size} + patch_size: ${patch_size} + debug: True + repeat_batch: True + dataset: + _target_: data.composed_dataset.ComposedDataset + dataset_configs: + - _target_: data.datasets.co3d.Co3dDataset + split: train + CO3D_DIR: /fsx-repligen/jianyuan/transfer_buffer/small_set/co3d/ + CO3D_ANNOTATION_DIR: /fsx-repligen/jianyuan/transfer_buffer/small_set/co3d_anno + val: + _target_: data.dynamic_dataloader.DynamicTorchDataset + num_workers: ${num_workers} + common_config: + img_size: ${img_size} + patch_size: ${patch_size} + debug: True + dataset: + _target_: data.composed_dataset.ComposedDataset + dataset_configs: + - _target_: data.datasets.co3d.Co3dDataset + split: test + CO3D_DIR: /fsx-repligen/jianyuan/transfer_buffer/small_set/co3d/ + CO3D_ANNOTATION_DIR: /fsx-repligen/jianyuan/transfer_buffer/small_set/co3d_anno + + +logging: + log_dir: logs + log_visuals: False + log_freq: 1 + log_level_primary: DEBUG + log_level_secondary: WARNING + all_ranks: False + tensorboard_writer: + _target_: train_utils.tb_writer.TensorBoardLogger + path: ${logging.log_dir}/tensorboard + scalar_keys_to_log: + train: + keys_to_log: + - loss_objective + - loss_camera + - loss_T + - loss_R + - loss_FL + - loss_conf_depth + - loss_reg_depth + - loss_grad_depth + val: + keys_to_log: + - loss_objective + - loss_camera + - loss_T + - loss_R + - loss_FL + - loss_conf_depth + - loss_reg_depth + - loss_grad_depth + + + +checkpoint: + save_dir: logs/${exp_name}/ckpts + save_freq: 5 + resume_checkpoint_path: /fsx-repligen/jianyuan/transfer_buffer/ckpts/model.pt + strict: False + + +loss: + _target_: loss.MultitaskLoss + camera: + weight: 5.0 + loss_type: "l1" + depth: + weight: 1.0 + gradient_loss_fn: "grad" + valid_range: 0.98 + point: null + track: null + + + + +optim: + param_group_modifiers: False + + optimizer: + _target_: torch.optim.AdamW + lr: 1e-4 + weight_decay: 0.05 + + # KEY CHANGE: No freezing — all parameters are trainable + frozen_module_names: [] + + amp: + enabled: True + amp_dtype: bfloat16 + gradient_clip: + _target_: train_utils.gradient_clip.GradientClipper + configs: + - module_name: ["aggregator"] + max_norm: 1.0 + norm_type: 2 + - module_name: ["depth"] + max_norm: 1.0 + norm_type: 2 + - module_name: ["camera"] + max_norm: 1.0 + norm_type: 2 + options: + lr: + - scheduler: + _target_: fvcore.common.param_scheduler.CompositeParamScheduler + schedulers: + - _target_: fvcore.common.param_scheduler.LinearParamScheduler + start_value: 1e-8 + end_value: 1e-4 + - _target_: fvcore.common.param_scheduler.CosineParamScheduler + start_value: 1e-4 + end_value: 1e-8 + lengths: [0.05, 0.95] + interval_scaling: ['rescaled', 'rescaled'] + weight_decay: + - scheduler: + _target_: fvcore.common.param_scheduler.ConstantParamScheduler + value: 0.05 + + + + +max_epochs: 100 + +model: + _target_: vggt.models.vggt.VGGT + enable_camera: True + enable_depth: True + enable_point: False + enable_track: False + + +distributed: + backend: nccl + comms_dtype: None + find_unused_parameters: False + timeout_mins: 30 + gradient_as_bucket_view: True + bucket_cap_mb: 25 + broadcast_buffers: True + +cuda: + cudnn_deterministic: False + cudnn_benchmark: False + allow_tf32: True diff --git a/training/launch.py b/training/launch.py index 2ad826c12..fb411dc39 100644 --- a/training/launch.py +++ b/training/launch.py @@ -8,5 +8,3 @@ trainer = Trainer(**cfg) trainer.run() -import pdb;pdb.set_trace() -m=1 diff --git a/training/launch_multi.py b/training/launch_multi.py new file mode 100644 index 000000000..5062b000c --- /dev/null +++ b/training/launch_multi.py @@ -0,0 +1,119 @@ +""" +灵活的 VGGT 训练启动器 - 支持多种实验配置 + +用法: + # 默认配置（部分参数微调，冻结 Aggregator） + torchrun --nproc_per_node=1 launch_multi.py --config default + + # 全参微调（需要 H100 或以上 GPU） + torchrun --nproc_per_node=1 launch_multi.py --config full_finetune + + # 消融实验 1: 移除深度损失 + torchrun --nproc_per_node=1 launch_multi.py --config ablation_no_depth + + # 消融实验 2: 移除相机损失 + torchrun --nproc_per_node=1 launch_multi.py --config ablation_no_camera + + # 消融实验 3: 移除深度梯度损失 + torchrun --nproc_per_node=1 launch_multi.py --config ablation_no_grad_depth + +多 GPU 训练: + torchrun --nproc_per_node=4 launch_multi.py --config default + +覆盖 Co3D 数据路径: + torchrun --nproc_per_node=1 launch_multi.py --config default \\ + --override data.train.dataset.dataset_configs.0.CO3D_DIR=/path/to/co3d \\ + --override data.train.dataset.dataset_configs.0.CO3D_ANNOTATION_DIR=/path/to/co3d_anno \\ + --override data.val.dataset.dataset_configs.0.CO3D_DIR=/path/to/co3d \\ + --override data.val.dataset.dataset_configs.0.CO3D_ANNOTATION_DIR=/path/to/co3d_anno +""" + +import os +import sys +import argparse + +from hydra import initialize, compose +from omegaconf import DictConfig, OmegaConf +from trainer import Trainer + + +def parse_args(): + parser = argparse.ArgumentParser(description="VGGT Training Launcher") + parser.add_argument("--config", type=str, default="default", + choices=["default", "full_finetune", + "ablation_no_depth", "ablation_no_camera", + "ablation_no_grad_depth"], + help="配置文件名称（不含 .yaml 后缀）") + parser.add_argument("--override", nargs="*", default=[], + help="覆盖配置项，格式: key=value") + return parser.parse_args() + + +def apply_overrides(cfg, overrides): + """应用命令行覆盖到配置""" + for override in overrides: + if "=" not in override: + print(f"WARNING: Invalid override format: {override}, expected key=value") + continue + key, value = override.split("=", 1) + # 尝试转换类型 + try: + if value.lower() == "true": + value = True + elif value.lower() == "false": + value = False + elif value.lower() == "null": + value = None + elif value.isdigit(): + value = int(value) + elif value.replace(".", "").replace("-", "").isdigit(): + value = float(value) + except (ValueError, AttributeError): + pass # keep as string + + OmegaConf.update(cfg, key, value, merge=True) + print(f"Override: {key} = {value}") + + return cfg + + +def main(): + args = parse_args() + config_name = args.config + + print("=" * 70) + print(f"VGGT Training Launcher") + print(f"Configuration: {config_name}.yaml") + print(f"Overrides: {args.override}") + print("=" * 70) + + with initialize(version_base=None, config_path="config"): + cfg = compose(config_name=config_name) + + # 应用命令行覆盖 + if args.override: + cfg = apply_overrides(cfg, args.override) + + # 打印关键配置信息 + print("\n--- 训练配置摘要 ---") + print(f"实验名称: {cfg.exp_name}") + print(f"图像大小: {cfg.img_size}") + print(f"最大 Epochs: {cfg.max_epochs}") + print(f"梯度累积步数: {cfg.accum_steps}") + print(f"学习率: {cfg.optim.optimizer.lr}") + print(f"冻结模块: {cfg.optim.frozen_module_names}") + print(f"Camera Loss: {cfg.loss.camera is not None}") + print(f"Depth Loss: {cfg.loss.depth is not None}") + if cfg.loss.depth: + print(f" Depth Gradient Loss: {cfg.loss.depth.gradient_loss_fn}") + print(f"Point Loss: {cfg.loss.point is not None}") + print(f"AMP: {cfg.optim.amp.enabled} ({cfg.optim.amp.amp_dtype})") + print("-" * 70) + + # 初始化并运行 Trainer + trainer = Trainer(**cfg) + trainer.run() + + +if __name__ == "__main__": + main() diff --git a/training/loss.py b/training/loss.py index f919c7c1a..fbeac9b99 100644 --- a/training/loss.py +++ b/training/loss.py @@ -316,13 +316,13 @@ def regression_loss(pred, gt, mask, conf=None, gradient_loss_fn=None, gamma=1.0, loss_grad = 0 # Prepare confidence for gradient loss if needed - if "conf" in gradient_loss_fn: + if gradient_loss_fn is not None and "conf" in gradient_loss_fn: to_feed_conf = conf.reshape(bb*ss, hh, ww) else: to_feed_conf = None # Compute gradient loss if specified for spatial smoothness - if "normal" in gradient_loss_fn: + if gradient_loss_fn is not None and "normal" in gradient_loss_fn: # Surface normal-based gradient loss loss_grad = gradient_loss_multi_scale_wrapper( pred.reshape(bb*ss, hh, ww, nc), @@ -332,7 +332,7 @@ def regression_loss(pred, gt, mask, conf=None, gradient_loss_fn=None, gamma=1.0, scales=3, conf=to_feed_conf, ) - elif "grad" in gradient_loss_fn: + elif gradient_loss_fn is not None and "grad" in gradient_loss_fn: # Standard gradient-based loss loss_grad = gradient_loss_multi_scale_wrapper( pred.reshape(bb*ss, hh, ww, nc), diff --git a/training/trainer.py b/training/trainer.py index 6577b95c6..d7ab12d44 100644 --- a/training/trainer.py +++ b/training/trainer.py @@ -537,7 +537,6 @@ def val_epoch(self, val_loader, is_fresh_epoch): batch = self._process_batch(batch, 'val', local_data_ids) batch = copy_data_to_device(batch, self.device) - import pdb; pdb.set_trace() # compute output with torch.no_grad(): with torch.cuda.amp.autocast(