Marigold开发者手册深入理解训练代码和自定义扩展【免费下载链接】Marigold[CVPR 2024 - Oral, Best Paper Award Candidate] Marigold: Repurposing Diffusion-Based Image Generators for Monocular Depth Estimation项目地址: https://gitcode.com/gh_mirrors/ma/MarigoldMarigold是一个革命性的计算机视觉项目它重新利用预训练的扩散图像生成器进行单目深度估计、表面法线预测和内在图像分解。这篇开发者手册将深入解析Marigold的训练代码架构、自定义扩展方法和最佳实践帮助你充分利用这个强大的框架。 Marigold训练架构概述Marigold的核心思想是利用预训练的Stable Diffusion模型中的丰富视觉知识通过微调将其重新用于密集图像分析任务。训练代码位于script/depth/train.py、script/normals/train.py和script/iid/train.py分别对应深度估计、表面法线和内在图像分解任务。训练流程核心组件配置系统基于OmegaConf的YAML配置文件如config/train_marigold_depth.yaml数据集架构基础数据集类位于src/dataset/base_depth_dataset.py训练器实现核心训练逻辑在src/trainer/marigold_depth_trainer.pyPipeline组件推理和训练管道在marigold/marigold_depth_pipeline.py️ 训练代码结构解析1. 训练配置详解Marigold使用模块化的配置文件系统。每个任务都有对应的训练配置文件深度估计config/train_marigold_depth.yaml表面法线config/train_marigold_normals.yaml内在图像分解config/train_marigold_iid_appearance.yaml配置文件采用分层设计通过base_config引用基础配置base_config: - config/logging.yaml - config/wandb.yaml - config/dataset_depth/dataset_train.yaml - config/dataset_depth/dataset_val.yaml - config/dataset_depth/dataset_vis.yaml - config/model_sdv2.yaml2. 数据集自定义指南Marigold支持多种数据集所有数据集都继承自基础类。要添加自定义数据集需要创建数据集类继承相应的基础类深度估计继承BaseDepthDataset表面法线继承BaseNormalsDataset内在图像分解继承BaseIIDDataset实现核心方法def __getitem__(self, index): # 加载RGB图像 # 加载对应的标注深度/法线/内在分解 # 应用数据增强 # 返回标准化后的数据注册数据集在src/dataset/init.py中添加导入3. 训练器工作原理Marigold训练器的核心在src/trainer/marigold_depth_trainer.py中实现。关键步骤包括模型适配修改U-Net的输入通道数优化器配置使用Adam优化器学习率3e-5损失函数默认使用MSE损失噪声调度使用DDPM调度器进行训练图Marigold深度估计流程展示 - 从原始图像到深度热力图再到3D点云 自定义扩展实战指南1. 添加新的训练任务要添加新的训练任务需要创建新的Pipeline类继承自MarigoldDepthPipeline实现训练器参考src/trainer/marigold_depth_trainer.py配置训练脚本创建新的训练脚本和配置文件2. 修改模型架构Marigold允许修改U-Net架构以适应不同任务。关键修改点# 在训练器中修改输入通道 def _replace_unet_conv_in(self): # 修改U-Net的第一层卷积以适应不同输入 conv_in Conv2d( 8, # 新输入通道数 self.model.unet.config[block_out_channels][0], kernel_size3, padding1, ) self.model.unet.conv_in conv_in3. 数据预处理自定义数据预处理配置在数据集配置文件中。例如深度估计的数据集配置位于config/dataset_depth/dataset_train.yamldataset: train: name: mixed prob_ls: [0.9, 0.1] dataset_list: - name: hypersim_depth disp_name: hypersim_depth_train dir: hypersim/hypersim_processed_train.tar filenames: data_split/hypersim_depth/filename_list_train_filtered.txt resize_to_hw: - 480 - 640 训练最佳实践1. 环境设置确保正确设置环境变量export BASE_DATA_DIRYOUR_DATA_DIR export BASE_CKPT_DIRYOUR_CHECKPOINT_DIR2. 数据准备不同任务需要不同的数据集深度估计Hypersim和Virtual KITTI 2表面法线Hypersim、Interiorverse和Sintel内在图像分解Interiorverse外观模型和Hypersim光照模型3. 训练命令启动训练的基本命令# 深度估计 python script/depth/train.py --config config/train_marigold_depth.yaml # 表面法线 python script/normals/train.py --config config/train_marigold_normals.yaml # 内在图像分解 python script/iid/train.py --config config/train_marigold_iid_appearance.yaml4. 恢复训练从检查点恢复训练python script/depth/train.py --resume_run output/marigold_base/checkpoint/latest 调试和优化技巧1. 验证和可视化训练过程中定期进行验证和可视化validation: denoising_steps: 1 ensemble_size: 1 processing_res: 0 match_input_res: false resample_method: bilinear2. 学习率调度使用指数衰减学习率调度lr_scheduler: name: IterExponential kwargs: total_iter: 25000 final_ratio: 0.01 warmup_steps: 1003. 多分辨率噪声增强训练稳定性的多分辨率噪声multi_res_noise: strength: 0.9 annealed: true downscale_strategy: original 评估和部署1. 模型评估使用评估脚本评估模型性能# 深度估计 bash script/depth/eval/11_infer_nyu.sh bash script/depth/eval/12_eval_nyu.sh2. 检查点组合训练后组合检查点将训练输出中的unet文件夹替换到Marigold检查点中替换scheduler/scheduler_config.json文件3. 集成到DiffusersMarigold已集成到Hugging Face Diffusers库可以直接使用from diffusers import MarigoldDepthPipeline pipeline MarigoldDepthPipeline.from_pretrained(prs-eth/marigold-depth-v1-1) 高级自定义技巧1. 自定义损失函数在src/util/loss.py中添加自定义损失函数def custom_loss(pred, target, **kwargs): # 实现自定义损失逻辑 return loss_value2. 数据增强策略在配置文件中调整数据增强参数augmentation: lr_flip_p: 0.5 # 添加更多增强策略3. 混合数据集采样使用MixedBatchSampler实现多数据集混合训练。 实用建议和注意事项硬件要求建议使用至少24GB显存的GPU进行训练训练时间完整训练通常需要2-3天数据格式确保数据格式与现有数据集保持一致版本兼容性注意PyTorch和Diffusers库的版本兼容性内存优化使用梯度累积处理大批量数据通过这篇开发者手册你应该对Marigold的训练代码架构和自定义扩展有了深入理解。无论是添加新数据集、修改模型架构还是创建新任务Marigold的模块化设计都为你提供了灵活的扩展能力。【免费下载链接】Marigold[CVPR 2024 - Oral, Best Paper Award Candidate] Marigold: Repurposing Diffusion-Based Image Generators for Monocular Depth Estimation项目地址: https://gitcode.com/gh_mirrors/ma/Marigold创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考