src/main.py

import os
from pathlib import Path
import warnings
import copy

import hydra
import torch
import swanlab as wandb
import os
from colorama import Fore
from jaxtyping import install_import_hook
from omegaconf import DictConfig, OmegaConf

import random
import numpy as np

import sys
# 排除冲突路径
sys.path = [p for p in sys.path if "/mnt/data-3/users/nichaojun/.local" not in p]


from pytorch_lightning import Trainer
from pytorch_lightning.callbacks import (
    LearningRateMonitor,
    ModelCheckpoint,
)

from pytorch_lightning.callbacks import Callback
from typing import Any, Dict, Optional


# from pytorch_lightning.loggers.wandb import WandbLogger
from swanlab.integration.pytorch_lightning import SwanLabLogger

from pytorch_lightning.plugins.environments import LightningEnvironment
from pytorch_lightning.strategies import DDPStrategy  

# Configure beartype and jaxtyping.
with install_import_hook(
    ("src",),
    ("beartype", "beartype"),
):
    from src.config import load_typed_root_config
    from src.dataset.data_module import DataModule
    from src.global_cfg import set_cfg
    from src.loss import get_losses
    from src.misc.LocalLogger import LocalLogger
    from src.misc.step_tracker import StepTracker
    from src.misc.wandb_tools import update_checkpoint_path
    from src.misc.resume_ckpt import find_latest_ckpt
    from src.model.decoder import get_decoder
    from src.model.encoder import get_encoder
    from src.model.model_wrapper import ModelWrapper


def cyan(text: str) -> str:
    return f"{Fore.CYAN}{text}{Fore.RESET}"


class UnfreezePretrainedCallback(Callback):
    def __init__(self, unfreeze_step: int = 20000):  # 修改1：参数名改为 unfreeze_step
        self.unfreeze_step = unfreeze_step  #
        self.has_unfrozen = False

    def on_train_batch_end(self, trainer, pl_module, outputs, batch, batch_idx):  # 修改3：改为 batch_end 钩子
        # 只在主进程执行解冻操作
        if trainer.is_global_zero and not self.has_unfrozen:
            current_step = trainer.global_step  # 修改4：获取当前步数
            if current_step >= self.unfreeze_step:
                print(cyan(f"Step {current_step}: Unfreezing pretrained_monodepth parameters"))
                # 解冻所有参数
                for param in pl_module.encoder.depth_predictor.parameters():
                    param.requires_grad = True

                self.has_unfrozen = True

def _set_global_seed(seed):
    """设置Python和NumPy的全局随机种子"""
    random.seed(seed)
    np.random.seed(seed)
    # 注意：这里不设置PyTorch的种子，因为PyTorch种子需要按进程单独设置
    print(f"Global Python/NumPy seed set to: {seed}")



@hydra.main(
    version_base=None,
    config_path="../config",
    config_name="main",
)

def train(cfg_dict: DictConfig):

    # # 设置GPU内存限制 - 在任何CUDA操作之前
    # if torch.cuda.is_available():
    #     # 你可以从配置中读取内存限制比例，或者直接设置一个固定值
    #     memory_fraction = cfg_dict.get("gpu_memory_fraction", 0.8)  # 默认使用80%显存
    #     torch.cuda.set_per_process_memory_fraction(memory_fraction)
    #     print(f"Set GPU memory fraction to: {memory_fraction}")


    if cfg_dict["mode"] == "train" and cfg_dict["train"]["eval_model_every_n_val"] > 0:
        eval_cfg_dict = copy.deepcopy(cfg_dict)
        dataset_dir = str(cfg_dict["dataset"]["roots"]).lower()
        if "re10k" in dataset_dir:
            eval_path = "assets/evaluation_index_re10k.json"
        elif "dl3dv" in dataset_dir:
            if cfg_dict["dataset"]["view_sampler"]["num_context_views"] == 6:
                eval_path = "assets/dl3dv_start_0_distance_50_ctx_6v_tgt_8v.json"
            else:
                raise ValueError("unsupported number of views for dl3dv")
        else:
            raise Exception("Fail to load eval index path")
        eval_cfg_dict["dataset"]["view_sampler"] = {
            "name": "evaluation",
            "index_path": eval_path,
            "num_context_views": cfg_dict["dataset"]["view_sampler"]["num_context_views"],
        }
        eval_cfg = load_typed_root_config(eval_cfg_dict)
    else:
        eval_cfg = None

    cfg = load_typed_root_config(cfg_dict)
    set_cfg(cfg_dict)

    # Set up the output directory.
    if cfg_dict.output_dir is None:
        output_dir = Path(
            hydra.core.hydra_config.HydraConfig.get()["runtime"]["output_dir"]
        )
    else:  # for resuming
        output_dir = Path(cfg_dict.output_dir)
        os.makedirs(output_dir, exist_ok=True)
    print(cyan(f"Saving outputs to {output_dir}."))

    # Set up logging with wandb.
    callbacks = []
    if cfg_dict.wandb.mode != "disabled" and cfg.mode == "train":
        wandb_extra_kwargs = {}
        if cfg_dict.wandb.id is not None:
            wandb_extra_kwargs.update({'id': cfg_dict.wandb.id,
                                       'resume': "must"})
        # logger = WandbLogger( 
        #     entity=cfg_dict.wandb.entity,
        #     project=cfg_dict.wandb.project,
        #     mode=cfg_dict.wandb.mode,
        #     name=os.path.basename(cfg_dict.output_dir),
        #     tags=cfg_dict.wandb.get("tags", None),
        #     log_model=False,
        #     save_dir=output_dir,
        #     config=OmegaConf.to_container(cfg_dict),
        #     **wandb_extra_kwargs,
        # )
        logger = SwanLabLogger(
            project = cfg_dict.wandb.project,
            experiment_name = cfg_dict.wandb.entity,
            workspace = cfg_dict.wandb.workspace
        )

        callbacks.append(LearningRateMonitor("step", True))

        # if wandb.run is not None:
        #     wandb.run.log_code("src")
    else:
        logger = LocalLogger()

    # Set up checkpointing.
    callbacks.append(
        ModelCheckpoint(
            output_dir / "checkpoints",
            every_n_train_steps=cfg.checkpointing.every_n_train_steps,
            save_top_k=cfg.checkpointing.save_top_k,
            monitor="info/global_step",
            mode="max",
        )
    )

    #解冻参数回调函数
    callbacks.append(UnfreezePretrainedCallback(unfreeze_step=20000))

    for cb in callbacks:
        cb.CHECKPOINT_EQUALS_CHAR = '_'

    # Prepare the checkpoint for loading.
    if cfg.checkpointing.resume:
        if not os.path.exists(output_dir / 'checkpoints'):
            checkpoint_path = None
        else:
            checkpoint_path = find_latest_ckpt(output_dir / 'checkpoints')
            print(f'resume from {checkpoint_path}')
    else:
        checkpoint_path = update_checkpoint_path(cfg.checkpointing.load, cfg.wandb)

    # This allows the current step to be shared with the data loader processes.
    step_tracker = StepTracker()

    # 创建分布式策略
    if torch.cuda.device_count() > 1:
        # 创建DDP策略并启用未使用参数检测
        ddp_strategy = DDPStrategy(
            find_unused_parameters=True,  # 关键设置：解决冻结参数导致的DDP问题
            static_graph=False,  # 设置为False以适应冻结参数
            process_group_backend="nccl" if torch.cuda.is_available() else "gloo",
        )
    else:
        ddp_strategy = "auto"

    trainer = Trainer(
        max_epochs=-1,
        accelerator="gpu",
        logger=logger,
        devices=torch.cuda.device_count(),
        strategy= ddp_strategy , #'ddp' if torch.cuda.device_count() > 1 else "auto"
        callbacks=callbacks,
        val_check_interval=cfg.trainer.val_check_interval,
        enable_progress_bar=cfg.mode == "test",
        gradient_clip_val=cfg.trainer.gradient_clip_val,
        max_steps=cfg.trainer.max_steps,
        num_sanity_val_steps=cfg.trainer.num_sanity_val_steps,
        num_nodes=cfg.trainer.num_nodes,
        plugins=LightningEnvironment() if cfg.use_plugins else None,
    )

     # 只在主进程中设置全局种子，避免分布式环境下的问题
    if torch.distributed.is_initialized():
        # 在分布式环境中，只有rank 0设置全局种子
        if torch.distributed.get_rank() == 0:
            _set_global_seed(42)
        # 等待rank 0完成种子设置
        torch.distributed.barrier()
    else:
        # 在非分布式环境中，直接设置
        _set_global_seed(42)

    torch.manual_seed(cfg_dict.seed + trainer.global_rank)

    encoder, encoder_visualizer = get_encoder(cfg.model.encoder)

    model_wrapper = ModelWrapper(
        cfg.optimizer,
        cfg.test,
        cfg.train,
        encoder,
        encoder_visualizer,
        get_decoder(cfg.model.decoder, cfg.dataset),
        get_losses(cfg.loss),
        step_tracker,
        eval_data_cfg=(
            None if eval_cfg is None else eval_cfg.dataset
        ),
    )
    #测试参数
    model_wrapper._check_param_updates = False
    
    
    data_module = DataModule(
        cfg.dataset,
        cfg.data_loader,
        step_tracker,
        global_rank=trainer.global_rank,
    )

    if cfg.mode == "train":
        print("train:", len(data_module.train_dataloader()))
        print("val:", len(data_module.val_dataloader()))
        print("test:", len(data_module.test_dataloader()))

    strict_load = not cfg.checkpointing.no_strict_load

    if cfg.mode == "train":
        # only load monodepth
        if cfg.checkpointing.pretrained_monodepth is not None:
            strict_load = False
            pretrained_model = torch.load(cfg.checkpointing.pretrained_monodepth, map_location='cpu')
            if 'state_dict' in pretrained_model:
                pretrained_model = pretrained_model['state_dict']

            load_result = model_wrapper.encoder.depth_predictor.load_state_dict(pretrained_model, strict=strict_load)
            
            # 获取成功加载的参数名称
            loaded_keys = set(pretrained_model.keys()) - set(load_result.unexpected_keys)
            
            # 精确冻结：仅冻结成功加载的参数
            for name, param in model_wrapper.encoder.depth_predictor.named_parameters():
                if name in loaded_keys:
                    param.requires_grad = False
            
            print("\n===== 参数加载报告 =====")
            print(f"✅ 成功加载参数数量: {len(loaded_keys)}")
            print(f"❄️ 冻结参数数量: {len(loaded_keys)}")
            print(f"⚠️ 缺失参数: {len(load_result.missing_keys)} 个")
            print(f"⚠️ 多余参数: {len(load_result.unexpected_keys)} 个")
            print(f"🛠️ 可训练参数数量: {len([p for p in model_wrapper.encoder.depth_predictor.parameters() if p.requires_grad])}")
            
            print(
                cyan(
                    f"Loaded pretrained monodepth (partial freezing): {cfg.checkpointing.pretrained_monodepth}"
                )
            )
        
        # load pretrained mvdepth
        if cfg.checkpointing.pretrained_mvdepth is not None:
            pretrained_model = torch.load(cfg.checkpointing.pretrained_mvdepth, map_location='cpu')['model']

            load_result =  model_wrapper.encoder.depth_predictor.load_state_dict(pretrained_model, strict=False)
            
            print("\n===== 参数加载报告 =====")
            print(f"✅ 成功加载参数数量: {len(pretrained_model) - len(load_result.unexpected_keys)}")
            print(f"⚠️ 缺失参数: {len(load_result.missing_keys)} 个")
            print(f"⚠️ 多余参数: {len(load_result.unexpected_keys)} 个")
            
            
            print(
                cyan(
                    f"Loaded pretrained mvdepth: {cfg.checkpointing.pretrained_mvdepth}"
                )
            )
        
        # load full model
        if cfg.checkpointing.pretrained_model is not None:
            strict_load = False
            pretrained_model = torch.load(cfg.checkpointing.pretrained_model, map_location='cpu')
            if 'state_dict' in pretrained_model:
                pretrained_model = pretrained_model['state_dict']

            model_wrapper.load_state_dict(pretrained_model, strict=strict_load)
            print(
                cyan(
                    f"Loaded pretrained weights: {cfg.checkpointing.pretrained_model}"
                )
            )

        # load pretrained depth
        if cfg.checkpointing.pretrained_depth is not None:
            
            strict_load = False
            pretrained_model = torch.load(cfg.checkpointing.pretrained_depth, map_location='cpu')
            if 'state_dict' in pretrained_model:
                pretrained_model = pretrained_model['state_dict']
            
            # pretrained_model = torch.load(cfg.checkpointing.pretrained_depth, map_location='cpu')['model']
            # strict_load = True
            load_result  = model_wrapper.encoder.depth_predictor.load_state_dict(pretrained_model, strict=strict_load)
            
            
            print("\n===== 参数加载报告 =====")
            print(f"✅ 成功加载参数数量: {len(pretrained_model) - len(load_result.unexpected_keys)}")
            print(f"⚠️ 缺失参数: {len(load_result.missing_keys)} 个")
            print(f"⚠️ 多余参数: {len(load_result.unexpected_keys)} 个")
            
            print(
                cyan(
                    f"Loaded pretrained depth: {cfg.checkpointing.pretrained_depth}"
                )
            )
            
        trainer.fit(model_wrapper, datamodule=data_module, ckpt_path=checkpoint_path)
    else:
        # load full model
        if cfg.checkpointing.pretrained_model is not None:
            pretrained_model = torch.load(cfg.checkpointing.pretrained_model, map_location='cpu')
            if 'state_dict' in pretrained_model:
                pretrained_model = pretrained_model['state_dict']

            model_wrapper.load_state_dict(pretrained_model, strict=strict_load)
            print(
                cyan(
                    f"Loaded pretrained weights: {cfg.checkpointing.pretrained_model}"
                )
            )

        # load pretrained depth model only
        if cfg.checkpointing.pretrained_depth is not None:
            pretrained_model = torch.load(cfg.checkpointing.pretrained_depth, map_location='cpu')['model']

            strict_load = True
            model_wrapper.encoder.depth_predictor.load_state_dict(pretrained_model, strict=strict_load)
            print(
                cyan(
                    f"Loaded pretrained depth: {cfg.checkpointing.pretrained_depth}"
                )
            )
            
        trainer.test(
            model_wrapper,
            datamodule=data_module,
            ckpt_path=checkpoint_path,
        )






if __name__ == "__main__":
    warnings.filterwarnings("ignore")
    torch.set_float32_matmul_precision('high')

    train()