models.py

import torchmetrics.classification
import torchmetrics.segmentation
import torchmetrics.utilities
import torchmetrics.utilities.data
from typing import Any
import torch
import torch.nn.functional as F
import lightning as L
import torch.nn as nn
from torch import optim
from utils import config
import torchmetrics
from pathlib import Path

from net.PytorchUNet.unet.unet_model import UNet3D
from data_pretrain import generate_DA_transfroms
from net.MedNeXt.nnunet_mednext import create_mednext_v1, MedNeXt
from net.Med3DMaxVit import Med3DMaxViT
from net.SwinUNet import SwinUNETR
from net.SwinUNERT_monai import SwinTransformer as SwinViT
from net.pretrain.PCRLv2.MICCAI_LITS2017.get_model import setup_pcrlv2_model
from net.research_contributions.DAE.BTCV_Finetune.main_for_ngc import get_dae_model
from net.pretrain.GVSL.PyTorch.Downstream.models_finetune.UNet_GVSL import UNet3D_GVSL

from thop import profile, clever_format


def setup_model(weight_path: str = None, model_name: str = None):
    """如果模型越来越多考虑使用 improtlib 模块进行导入，现在懒得管

    Returns:
        _type_: _description_
    """
    if model_name is not None:
        config.model_name = model_name
    print(f"load model {config.model_name}")
    if config.model_name == "native_3dunet":
        model = UNet3D()
    elif config.model_name == "MedNeXt":
        model = create_mednext_v1(
            num_input_channels=1,
            num_classes=config.num_classes,
            model_id=config.MedNeXt_model_id,  # S, B, M and L are valid model ids
            kernel_size=3,  # 3x3x3 and 5x5x5 were tested in publication
            deep_supervision=False,  # was used in publication
        )
    elif config.model_name == "Med3DMaxViT":
        model = Med3DMaxViT(config.num_classes)
    elif config.model_name == "SwinUNETR":
        vit = SwinViT(
            in_chans=config.vit_in_channels,
            embed_dim=config.vit_embed_dim,
            window_size=config.vit_window_size,
            patch_size=config.vit_patch_size,
            depths=config.vit_depths,
            num_heads=config.vit_heads,
            mlp_ratio=config.vit_mlp_ratio,
            qkv_bias=config.vit_qkv_bias,
            drop_rate=config.vit_drop_rate,
            attn_drop_rate=config.vit_attn_drop_rate,
            drop_path_rate=config.vit_dropout_path_rate,
            use_checkpoint=config.vit_use_checkpoint,
            spatial_dims=config.vit_spatial_dims,
            use_v2=config.vit_use_v2,
        )
        model = SwinUNETR(
            img_size=config.patch_shape,
            in_channels=1,
            out_channels=config.num_classes,
            feature_size=config.vit_embed_dim,
            vit=vit,
        )
    elif config.model_name == "SwinViT":
        model = SwinViT(
            in_chans=config.vit_in_channels,
            embed_dim=config.vit_embed_dim,
            window_size=config.vit_window_size,
            patch_size=config.vit_patch_size,
            depths=config.vit_depths,
            num_heads=config.vit_heads,
            mlp_ratio=config.vit_mlp_ratio,
            qkv_bias=config.vit_qkv_bias,
            drop_rate=config.vit_drop_rate,
            attn_drop_rate=config.vit_attn_drop_rate,
            drop_path_rate=config.vit_dropout_path_rate,
            use_checkpoint=config.vit_use_checkpoint,
            spatial_dims=config.vit_spatial_dims,
            use_v2=config.vit_use_v2,
        )
    elif config.model_name == "PCRLv2":
        model = setup_pcrlv2_model(weight_path=weight_path)
    elif config.model_name == "DAE":
        model = get_dae_model()
    elif config.model_name == "GVSL":
        model = UNet3D_GVSL(
            n_channels=1,
            n_classes=config.num_classes,
            pretrain_weight_path=Path(config.tensorboard_dir)
            / "AMOS"
            / "pretrain_AMOS"
            / "pretrain_GVSL"
            / "GVSL_epoch_1000.pth",
        )
    else:
        raise ValueError(f"model {config.model_name} not found")

    print(
        f"number of model parameters is {sum([i.numel() for i in model.parameters()])}"
    )
    print(f"loadding {config.model_name}")
    return model


def mednext_test():
    from net.MedNeXt.nnunet_mednext import create_mednext_v1

    model = create_mednext_v1(
        num_input_channels=1,
        num_classes=16,
        model_id="S",  # S, B, M and L are valid model ids
        kernel_size=3,  # 3x3x3 and 5x5x5 were tested in publication
        deep_supervision=False,  # was used in publication
    )

    model.to(config.gpu_id)

    input = torch.rand(size=(1, 1, 80, 80, 80)).to(config.gpu_id)
    logit = model(input)
    print(logit.max())
    print(logit.min())
    print(logit.shape)

    predit = nn.Softmax(dim=1)(logit)
    print(predit.shape)
    print(predit.max())
    print(predit.min())


def simple_test():
    # from net.SwinUNERT_monai import SwinUNETR

    model = setup_model()
    print(model)
    # model.to(config.gpu_id)
    # input = torch.rand(size=(1, 1, 96, 96, 96)).to(config.gpu_id)
    input = torch.rand(size=(2, 1, 64, 64, 64))
    print(model(input).shape)


def print_param():
    model: nn.Module = setup_model()
    for n, m in model.named_modules():
        print(n, sum(mm.numel() for mm in m.parameters()))


def thop_model():
    model = setup_model()
    input = torch.rand(size=(2, 1, 64, 64, 64))
    macs, params = profile(model, inputs=(input,))
    macs, params = clever_format([macs, params], "%.3f")
    print(macs)
    print(params)


if __name__ == "__main__":
    # simple_test()
    print_param()
    # thop_model()