main.py

from data_pretrain import (
    setup_data,
    setup_pretraining_data,
)

from data import get_loader
from models import setup_model
from torch import nn, optim
from utils import config, save_hparams, seed_all
from loss import (
    ce_loss,
    dc_medpy,
    dc_medpy_new,
    dc_po_medpy,
    jc_medpy,
    jc_po_medpy,
    hd_po_medpy,
    asd_po_medpy,
)
from models import setup_model
from torch.optim.lr_scheduler import ExponentialLR
from monai.optimizers import WarmupCosineSchedule
from trainers import Trainer
import datetime
import torch
from torch.distributed import init_process_group, destroy_process_group
import os
from monai.losses import DiceCELoss  # 使用 monai 的 DiceCELoss 作为损失函数
from pathlib import Path
from monai.inferers import SlidingWindowInferer
from tqdm import tqdm
from itertools import islice
from gental_man import occupy_gpu_memory, release_gpu_memory


def main_singlegpu(log_name: str = None, mode: str = "normal", weight_path: str = None):
    """单GPU训练

    Args:
        log_name (str, optional): 日志文件夹名称. Defaults to None.
        mode (str, optional): 实验性质, one of [normal, pretrain, downstream]. Defaults to "normal".
        weight_path (str, optional): 加载权重路径. Defaults to None.
    """
    # 防止忘记
    config.trigger_pretraining = False

    # 保存 yaml 文件
    save_hparams(log_name, mode)
    gpu_id = config.gpu_id
    model = setup_model(weight_path)
    # train_loder, val_loader = setup_data(dir=config.dataset_patch, only_train=False)
    print(f"已提前占用显存")
    gi = int(os.environ["LOCAL_RANK"]) if config.toggle_mulitgpu else gpu_id
    occupy_gpu_memory(gpu_ids=[gi], memory_ratios=[0.8])
    train_loader, val_loader, _ = get_loader()
    release_gpu_memory(gpu_ids=[gi])
    print(f"释放显存，准备训练")
    optimizer = optim.AdamW(model.parameters(), lr=config.start_up_lr)
    # schedular = ExponentialLR(optimizer, gamma=config.exp_lr_gamma)
    schedular = WarmupCosineSchedule(
        optimizer,
        warmup_steps=config.lr_warmup_epochs,
        t_total=config.train_epochs,
        warmup_multiplier=0.1,
    )
    loss_fn = {
        "main_loss": DiceCELoss(
            include_background=False, to_onehot_y=True, softmax=True
        )
    }

    metric_fn = {
        "val_dc": dc_medpy_new,
        "val_jc": jc_medpy,
        # "val_hd": hd_medpy,   由于计算复杂度太高，因此暂不考虑
        "val_ce_loss": ce_loss,
    }

    trainer = Trainer(
        model=model,
        loss_fn=loss_fn,
        metric_fn=metric_fn,
        optimizer=optimizer,
        schedular=schedular,
        train_loader=train_loader,
        val_loader=val_loader,
        log_prefix=log_name,
        gpu_id=gpu_id,
        weight_path=weight_path,
        mode=mode,
    )

    trainer.train()


def ddp_setup():
    torch.cuda.set_device(int(os.environ["LOCAL_RANK"]))
    init_process_group(backend="nccl")


def main_multigpu(log_name=None, mode: str = "normal", weight_path=None):
    """多GPU运行入口，在命令行使用一下命令启动：
    torchrun --standalone --nproc_per_node=[number of device_ids] main.py

    不过，再次运行单GPU时，可能需要检查一下 CUDA_VISIBLE_DEVICES 是否正确设置。

    Args:
        log_name (_type_, optional): _description_. Defaults to None.
    """
    # 免得忘记
    config.toggle_mulitgpu = True
    config.trigger_pretraining = False
    config.pretrain_batch_num = 3
    # config.model_name = "SwinUNETR"
    os.environ["CUDA_VISIBLE_DEVICES"] = str(config.device_ids)[1:-1]
    ddp_setup()
    main_singlegpu(log_name, mode, weight_path=weight_path)
    destroy_process_group()


def pretrain_singlegpu(log_name: str = None, mode: str = "pretrain"):
    # 以防万一忘了配置。
    config.trigger_pretraining = True

    save_hparams(log_name, mode)
    gpu_id = config.gpu_id
    model = setup_model()

    gi = int(os.environ["LOCAL_RANK"]) if config.toggle_mulitgpu else gpu_id
    occupy_gpu_memory(gpu_ids=[gi], memory_ratios=[0.5])
    train_loader, val_loader = setup_pretraining_data()
    release_gpu_memory(gpu_ids=[gi])
    optimizer = optim.AdamW(model.parameters(), lr=config.start_up_lr)
    # schedular = ExponentialLR(optimizer, gamma=0.9)
    schedular = WarmupCosineSchedule(
        optimizer,
        warmup_steps=config.lr_warmup_epochs,
        t_total=config.train_epochs,
        warmup_multiplier=0.1,
    )

    trainer = Trainer(
        model=model,
        optimizer=optimizer,
        schedular=schedular,
        log_prefix=log_name,
        train_pretrain_loader=train_loader,
        val_pretrain_loader=val_loader,
        gpu_id=gpu_id,
        mode=mode,
    )

    trainer.pretrain()


def pretrain_multigpu(log_name=None, mode: str = "pretrain"):
    """多GPU运行入口，在命令行使用一下命令启动：
    torchrun --standalone --nproc_per_node=[number of device_ids] main.py
    """
    # 免得忘记
    config.toggle_mulitgpu = True
    os.environ["CUDA_VISIBLE_DEVICES"] = str(config.device_ids)[1:-1]
    ddp_setup()
    pretrain_singlegpu(log_name, mode)
    destroy_process_group()


def inference(weight_path: str = None):

    # 修改配置防止忘记
    config.trigger_pretraining = False
    config.toggle_mulitgpu = False
    config.batch_num = 3
    config.pretrain_batch_num = 3
    config.test_batch_num = None
    # config.model_name = "SwinUNETR"

    model = setup_model().to(config.gpu_id)
    # model = setup_model()
    weight = torch.load(
        weight_path,
        map_location=f"cuda:{config.gpu_id}",
        # map_location=f"cpu",
    )["MODEL_STATE"]
    model.load_state_dict(weight)

    metric_fn = {
        # "test_dc": dc_medpy_new,
        "test_dc_po": dc_po_medpy,
        "test_jc_po": jc_po_medpy,
        # "test_hd_po": hd_po_medpy,
        "test_asd_po": asd_po_medpy,
        # "val_hd": hd_medpy,   由于计算复杂度太高，因此暂不考虑
        # "test_ce_loss": ce_loss,
    }
    metric_epoch = {i: [] for i in metric_fn.keys()}
    inferer = SlidingWindowInferer(
        roi_size=config.patch_shape,
        sw_batch_size=config.sw_batch_size,
        overlap=0.25,
        mode="constant",
    )
    occupy_gpu_memory(gpu_ids=[config.gpu_id], memory_ratios=[0.5])
    t, v, test_loader = get_loader()
    release_gpu_memory(gpu_ids=[config.gpu_id])

    pbar = tqdm(islice(test_loader, config.test_batch_num), desc="inferening...")

    for batch in pbar:
        image = batch["image"].to(config.gpu_id)
        label = batch["label"].to(config.gpu_id)
        # image = batch["image"]
        # label = batch["label"]
        with torch.no_grad():
            pred = inferer(image, model)
        metric_step = {i: func(pred, label) for i, func in metric_fn.items()}
        for i in metric_fn.keys():
            # if isinstance(metric_step[i], list):

            metric_epoch[i].append(metric_step[i])

    metric_epoch = {
        i: torch.nanmean(torch.tensor(j), dim=0) for i, j in metric_epoch.items()
    }
    print(metric_epoch)


if __name__ == "__main__":
    seed_all(config.all_seed)
    # main_singlegpu(log_name="AMOS", mode="normal")
    # main_multigpu(log_name="p_swinunetr_d_SwinUNETR_AMOS_p96")
    # pretrain_singlegpu(log_name="pretrain_custom_SwinUNETR_v3", mode="pretrain")

    # main_multigpu(
    #     log_name="p_v118_seed_0_d_SwinUNETR_AMOS",
    #     mode="downstream",
    #     weight_path="/public/cjh/workspace/AbdominalSegmentation/tensorboard_log/pretrain/pretrain_v118_SwinUNETR_AMOS_seed_0/last_model.pth",
    # )

    # main_singlegpu(
    #     log_name="p_swinunetr_d_SwinUNETR_AMOS",
    #     mode="downstream",
    #     weight_path="/public1/cjh/workspace/AbdominalSegmentation/tensorboard_log/pretrain/pretrain_swinunetr_SwinUNETR/model_bestValRMSE.pt",
    # )

    # main_multigpu(
    #     log_name="p_swinmm_d_SwinUNETR_AMOS",
    #     mode="downstream",
    #     weight_path="/public/cjh/workspace/AbdominalSegmentation/tensorboard_log/AMOS/pretrain_AMOS/pretrain_swinmm_SwinUNETR/testfinal_model.pth",
    # )

    # main_multigpu(
    #     log_name="p_gvsl_d_SwinUNETR_RAOS",
    #     mode="downstream",
    #     # weight_path="/public1/cjh/workspace/AbdominalSegmentation/tensorboard_log/pretrain/pretrain_GVSL_RAOS/GVSL_epoch_1000.pth",
    # )

    # main_singlegpu(
    #     log_name="BTCV",
    #     mode="normal",
    # )

    # pretrain_multigpu(log_name="pretrain_v118_SwinUNETR_AMOS_seed_0", mode="pretrain")

    # pretrain_multigpu(log_name="pretrain_custom_SwinUNETR_v14", mode="pretrain")

    # main_singlegpu(
    #     log_name="p_pcrl_d_pcrl_AMOS",
    #     mode="downstream",
    #     weight_path="/public1/cjh/workspace/AbdominalSegmentation/tensorboard_log/pretrain/pretrain_pcrlv2/pcrlv2_luna_pretask_0.8_100.pt",
    # )

    # main_multigpu(
    #     log_name="p_voco_seed_42_d_SwinUNETR_AMOS",
    #     mode="downstream",
    #     weight_path="/public/cjh/workspace/AbdominalSegmentation/tensorboard_log/AMOS/pretrain_AMOS/pretrain_VoCo_old/model_step250000.pt",
    # )

    # main_multigpu(
    #     log_name="p_dae_seed_42_d_SwinUNETR_AMOS",
    #     mode="downstream",
    #     weight_path="/public/cjh/workspace/AbdominalSegmentation/tensorboard_log/AMOS/pretrain_AMOS/pretrain_dae/ckpt_epoch_299.pth",
    # )

    # main_multigpu(
    #     log_name="p_gmim_20000_d_SwinUNETR_AMOS",
    #     mode="downstream",
    #     weight_path="/public/cjh/workspace/AbdominalSegmentation/tensorboard_log/AMOS/pretrain_AMOS/pretrain_gmim_AMOS/20000_ckpt.pt",
    # )

    inference(
        weight_path="/public/cjh/workspace/AbdominalSegmentation/tensorboard_log/downstream/p_v113_wo_rec_seed_0_d_SwinUNETR_AMOS/last_model.pth"
    )