Jupitern52/TextBraTS

1# Copyright 2020 - 2022 MONAI Consortium2# Licensed under the Apache License, Version 2.0 (the "License");3# you may not use this file except in compliance with the License.4# You may obtain a copy of the License at5#     http://www.apache.org/licenses/LICENSE-2.06# Unless required by applicable law or agreed to in writing, software7# distributed under the License is distributed on an "AS IS" BASIS,8# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.9# See the License for the specific language governing permissions and10# limitations under the License.11import warnings12warnings.filterwarnings("ignore", category=FutureWarning, module="transformers.utils.generic")13 14import argparse15import os16import numpy as np17import torch18import torch.distributed as dist19import torch.multiprocessing as mp20import torch.nn.parallel21import torch.utils.data.distributed22from optimizers.lr_scheduler import LinearWarmupCosineAnnealingLR23from trainer import run_training24from utils.data_utils import get_loader25from monai.losses import DiceLoss26from monai.metrics import DiceMetric27from utils.textswin_unetr import TextSwinUNETR28from monai.transforms import Activations, AsDiscrete, Compose29from monai.utils.enums import MetricReduction30import random31 32 33parser = argparse.ArgumentParser(description="TextBraTS segmentation pipeline for TextBRATS image-text dataset")34parser.add_argument("--checkpoint", default=None, help="start training from saved checkpoint")35parser.add_argument("--logdir", default="TextBraTS", type=str, help="directory to save the tensorboard logs")36parser.add_argument("--fold", default=0, type=int, help="data fold, 0 for validation and 1 for training")37parser.add_argument("--pretrained_model_name", default="model.pt", type=str, help="pretrained model name")38parser.add_argument("--data_dir", default="./data/TextBraTSData", type=str, help="dataset directory")39parser.add_argument("--json_list", default="./Train.json", type=str, help="dataset json file")40parser.add_argument("--save_checkpoint", action="store_true", help="save checkpoint during training")41parser.add_argument("--max_epochs", default=200, type=int, help="max number of training epochs")42parser.add_argument("--batch_size", default=2, type=int, help="number of batch size")43parser.add_argument("--sw_batch_size", default=4, type=int, help="number of sliding window batch size")44parser.add_argument("--optim_lr", default=1e-4, type=float, help="optimization learning rate")45parser.add_argument("--optim_name", default="adamw", type=str, help="optimization algorithm")46parser.add_argument("--reg_weight", default=1e-5, type=float, help="regularization weight")47parser.add_argument("--momentum", default=0.99, type=float, help="momentum")48parser.add_argument("--noamp", action="store_true", help="do NOT use amp for training")49parser.add_argument("--val_every", default=1, type=int, help="validation frequency")50parser.add_argument("--distributed", action="store_true", help="start distributed training")51parser.add_argument("--world_size", default=1, type=int, help="number of nodes for distributed training")52parser.add_argument("--rank", default=0, type=int, help="node rank for distributed training")53parser.add_argument("--dist-url", default="tcp://127.0.0.1:23456", type=str, help="distributed url")54parser.add_argument("--dist-backend", default="nccl", type=str, help="distributed backend")55parser.add_argument("--norm_name", default="instance", type=str, help="normalization name")56parser.add_argument("--workers", default=8, type=int, help="number of workers")57parser.add_argument("--feature_size", default=48, type=int, help="feature size")58parser.add_argument("--in_channels", default=4, type=int, help="number of input channels")59parser.add_argument("--out_channels", default=3, type=int, help="number of output channels")60parser.add_argument("--cache_dataset", action="store_true", help="use monai Dataset class")61parser.add_argument("--a_min", default=-175.0, type=float, help="a_min in ScaleIntensityRanged")62parser.add_argument("--a_max", default=250.0, type=float, help="a_max in ScaleIntensityRanged")63parser.add_argument("--b_min", default=0.0, type=float, help="b_min in ScaleIntensityRanged")64parser.add_argument("--b_max", default=1.0, type=float, help="b_max in ScaleIntensityRanged")65parser.add_argument("--space_x", default=1.5, type=float, help="spacing in x direction")66parser.add_argument("--space_y", default=1.5, type=float, help="spacing in y direction")67parser.add_argument("--space_z", default=2.0, type=float, help="spacing in z direction")68parser.add_argument("--roi_x", default=128, type=int, help="roi size in x direction")69parser.add_argument("--roi_y", default=128, type=int, help="roi size in y direction")70parser.add_argument("--roi_z", default=128, type=int, help="roi size in z direction")71parser.add_argument("--dropout_rate", default=0.0, type=float, help="dropout rate")72parser.add_argument("--dropout_path_rate", default=0.0, type=float, help="drop path rate")73parser.add_argument("--RandScaleIntensityd_prob", default=0.1, type=float, help="RandScaleIntensityd aug probability")74parser.add_argument("--RandShiftIntensityd_prob", default=0.1, type=float, help="RandShiftIntensityd aug probability")75parser.add_argument("--infer_overlap", default=0.5, type=float, help="sliding window inference overlap")76parser.add_argument("--lrschedule", default="warmup_cosine", type=str, help="type of learning rate scheduler")77parser.add_argument("--warmup_epochs", default=50, type=int, help="number of warmup epochs")78parser.add_argument("--resume_ckpt", action="store_true", help="resume training from pretrained checkpoint")79parser.add_argument("--smooth_dr", default=1e-6, type=float, help="constant added to dice denominator to avoid nan")80parser.add_argument("--smooth_nr", default=0.0, type=float, help="constant added to dice numerator to avoid zero")81parser.add_argument("--use_checkpoint", action="store_true", help="use gradient checkpointing to save memory")82parser.add_argument("--spatial_dims", default=3, type=int, help="spatial dimension of input data")83parser.add_argument("--use_ssl_pretrained", action="store_true", help="use SSL pretrained ckpt")84parser.add_argument(85    "--pretrained_dir",86    default="./runs/TextBraTS/",87    type=str,88    help="pretrained checkpoint directory",89)90parser.add_argument("--squared_dice", action="store_true", help="use squared Dice")91parser.add_argument("--seed", type=int, default=23,help="use random seed")92 93 94def main():95    args = parser.parse_args()96    args.amp = not args.noamp97    args.logdir = "./runs/" + args.logdir98    random.seed(args.seed)99    np.random.seed(args.seed)100    torch.manual_seed(args.seed)101    if args.distributed:102        torch.cuda.manual_seed_all(args.seed)103        args.ngpus_per_node = torch.cuda.device_count()104        print("Found total gpus", args.ngpus_per_node)105        args.world_size = args.ngpus_per_node * args.world_size106        mp.spawn(main_worker, nprocs=args.ngpus_per_node, args=(args,))107    else:108        torch.cuda.manual_seed(args.seed)109        main_worker(gpu=0, args=args)110 111 112def main_worker(gpu, args):113    if args.distributed:114        torch.multiprocessing.set_start_method("fork", force=True)115    np.set_printoptions(formatter={"float": "{: 0.3f}".format}, suppress=True)116    args.gpu = gpu117    if args.distributed:118        args.rank = args.rank * args.ngpus_per_node + gpu119        dist.init_process_group(120            backend=args.dist_backend, init_method=args.dist_url, world_size=args.world_size, rank=args.rank121        )122    torch.cuda.set_device(args.gpu)123    torch.backends.cudnn.benchmark = True124    args.test_mode = False125    loader = get_loader(args)126    print(args.rank, " gpu", args.gpu)127    if args.rank == 0:128        print("Batch size is:", args.batch_size, "epochs", args.max_epochs)129    pretrained_dir = args.pretrained_dir130    model_name = args.pretrained_model_name131    pretrained_pth = os.path.join(pretrained_dir, model_name)132 133    model = TextSwinUNETR(134        img_size=(args.roi_x, args.roi_y, args.roi_z),135        in_channels=args.in_channels,136        out_channels=args.out_channels,137        feature_size=args.feature_size,138        use_checkpoint=args.use_checkpoint,139        text_dim=768,140    )141 142    if args.resume_ckpt:143        model_dict = torch.load(pretrained_pth)["state_dict"]144        for key in list(model_dict.keys()):145            model_dict[key.replace("module.", "")] = model_dict.pop(key)146        model.load_state_dict(model_dict,strict=True)147        print("Using pretrained weights")148 149    if args.use_ssl_pretrained:150        try:151            model_dict = torch.load("/media/iipl/disk1/swinunetr/model_swinvit.pt",weights_only=True)152            state_dict = model_dict["state_dict"]153            # fix potential differences in state dict keys from pre-training to154            # fine-tuning155            for key in list(state_dict.keys()):156                state_dict[key.replace("module.", "swinViT.")] = state_dict.pop(key)157            for key in list(state_dict.keys()):158                if "fc" in key:159                    state_dict[key.replace("fc","linear")] = state_dict.pop(key)160                if "patch_embed" in key:161                    state_dict[key.replace("patch_embed","")] = state_dict.pop(key)162            model.load_state_dict(state_dict, strict=False)163        except ValueError:164            raise ValueError("Self-supervised pre-trained weights not available for" + str(args.model_name))165 166    if args.squared_dice:167        dice_loss = DiceLoss(168            to_onehot_y=False, sigmoid=True, squared_pred=True, smooth_nr=args.smooth_nr, smooth_dr=args.smooth_dr169        )170    else:171        dice_loss = DiceLoss(to_onehot_y=False, sigmoid=True)172    post_sigmoid = Activations(sigmoid=True)173    post_pred = AsDiscrete(argmax=False, logit_thresh=0.5)174    dice_acc = DiceMetric(include_background=True, reduction=MetricReduction.MEAN_BATCH, get_not_nans=True)175    pytorch_total_params = sum(p.numel() for p in model.parameters() if p.requires_grad)176    print("Total parameters count", pytorch_total_params)177 178    best_acc = 0179    start_epoch = 0180 181    if args.checkpoint is not None:182        checkpoint = torch.load(args.checkpoint, map_location="cpu")183        from collections import OrderedDict184 185        new_state_dict = OrderedDict()186        for k, v in checkpoint["state_dict"].items():187            new_state_dict[k.replace("backbone.", "")] = v188        model.load_state_dict(new_state_dict, strict=False)189        if "epoch" in checkpoint:190            start_epoch = checkpoint["epoch"]191        if "best_acc" in checkpoint:192            best_acc = checkpoint["best_acc"]193        print("=> loaded checkpoint '{}' (epoch {}) (bestacc {})".format(args.checkpoint, start_epoch, best_acc))194 195    model.cuda(args.gpu)196 197    if args.distributed:198        torch.cuda.set_device(args.gpu)199        if args.norm_name == "batch":200            model = torch.nn.SyncBatchNorm.convert_sync_batchnorm(model)201        model.cuda(args.gpu)202        model = torch.nn.parallel.DistributedDataParallel(model, device_ids=[args.gpu], output_device=args.gpu, find_unused_parameters = False,)203    if args.optim_name == "adam":204        optimizer = torch.optim.Adam(model.parameters(), lr=args.optim_lr, weight_decay=args.reg_weight)205    elif args.optim_name == "adamw":206        optimizer = torch.optim.AdamW(model.parameters(), lr=args.optim_lr, weight_decay=args.reg_weight)207    elif args.optim_name == "sgd":208        optimizer = torch.optim.SGD(209            model.parameters(), lr=args.optim_lr, momentum=args.momentum, nesterov=True, weight_decay=args.reg_weight210        )211    else:212        raise ValueError("Unsupported Optimization Procedure: " + str(args.optim_name))213 214    if args.lrschedule == "warmup_cosine":215        scheduler = LinearWarmupCosineAnnealingLR(216            optimizer, warmup_epochs=args.warmup_epochs, max_epochs=args.max_epochs217        )218    elif args.lrschedule == "cosine_anneal":219        scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=args.max_epochs)220        if args.checkpoint is not None:221            scheduler.step(epoch=start_epoch)222    else:223        scheduler = None224 225    semantic_classes = ["Dice_Val_TC", "Dice_Val_WT", "Dice_Val_ET"]226 227    accuracy = run_training(228        model=model,229        train_loader=loader[0],230        val_loader=loader[1],231        optimizer=optimizer,232        loss_func=dice_loss,233        acc_func=dice_acc,234        args=args,235        scheduler=scheduler,236        start_epoch=start_epoch,237        post_sigmoid=post_sigmoid,238        post_pred=post_pred,239        semantic_classes=semantic_classes,240    )241    return accuracy242 243 244if __name__ == "__main__":245    main()246