OUT OF MEMORY ENSEMBLE EVALUATION

[casarinsof]

Hello, I am having trouble with ensemble evaluator of MONAI, as even if the training went well, when using ensemble evaluator I get a CUDA out of memory error. Is there a way to limit memory usage? It seems that after predicting each pred_i of each fold (after a k fold) it still has the model loaded when going to the next model

[KumoLiu]

Hi @casarinsof,
I would suggest you to find some tips from this ensemble tutorial https://github.com/Project-MONAI/tutorials/blob/main/modules/cross_validation_models_ensemble.ipynb. If you still have some problem, feel free to ask me!

[casarinsof]

Hello, thanks for the reply.
I have built my code based on that tutorial, to be more precise it is basically what I find there with the exception of course for the model. In my case it's a Dynamic Unet (of Monai).
Still I cannot figure out how to save memory while using ensemble evaluation. Any tips?
Thanks

[KumoLiu]

Maybe you can use amp=True in EnsembleEvaluator to reduce memory usage and also reduce the number of your batch size.

[casarinsof]

Ok thanks, I will try with amp=True. For the batch, I am using already batch=1.

My code is:

def ensemble_evaluate(post_transforms, models, device, loader, inferer, metrics_saver):
evaluator = EnsembleEvaluator(
device=device,
val_data_loader=loader,
pred_keys=["pred0", "pred1", "pred2", "pred3", "pred4"],
networks=models,
inferer=inferer,
postprocessing=post_transforms,
val_handlers=metrics_saver,
key_val_metric={
"test_mean_dice": MeanDice(
include_background=True,
output_transform=from_engine(["pred", "label"])
),
},
)
evaluator.run()

def main():
args = parser.parse_args()
args.test_mode = True
output_directory = './outputs/'+args.exp_name
if not os.path.exists(output_directory):
os.makedirs(output_directory)
test_loader = get_loader(args)
pretrained_dir = args.pretrained_dir
inf_size = [args.roi_x, args.roi_y, args.roi_z]
inf_space = [args.space_x, args.space_y, args.space_z]

num = args.nfolds
models = []
for folder in range(num):
    print(f"--------------------Folder {folder}-------------------\n")

    model_name = f'{args.pretrained_model_name}{folder}.pt'
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    pretrained_pth = os.path.join(pretrained_dir, model_name)
    if (args.model_name is None) or args.model_name == "unetr":
        model = UNETR(
            in_channels=args.in_channels,
            out_channels=args.out_channels,
            img_size=(args.roi_x, args.roi_y, args.roi_z),
            feature_size=args.feature_size,
            hidden_size=args.hidden_size,
            mlp_dim=args.mlp_dim,
            num_heads=args.num_heads,
            pos_embed=args.pos_embed,
            norm_name=args.norm_name,
            conv_block=True,
            res_block=True,
            dropout_rate=args.dropout_rate,
        )

    elif args.model_name == 'swin_unetr':
        model = SwinUNETR(img_size=(args.roi_x, args.roi_y, args.roi_z),
                          in_channels=args.in_channels,
                          out_channels=args.out_channels,
                          feature_size=args.feature_size,
                          drop_rate=0.0,
                          attn_drop_rate=0.0,
                          dropout_path_rate=args.dropout_path_rate,
                          use_checkpoint=args.use_checkpoint,
                          )

    elif args.model_name == 'nnunet':
        kernels, strides = get_kernels_strides(inf_size, inf_space)
        if args.kernels is not None:
            kernels = args.kernels
        if args.strides is not None:
            strides = args.strides
        model = DynUNet(
            spatial_dims=args.spatial_dims,
            in_channels=args.in_channels,
            out_channels=args.out_channels,
            kernel_size=kernels,
            strides=strides,
            upsample_kernel_size=strides[1:],
            filters=args.filters,
            dropout=args.dropout_rate,
            norm_name=args.norm_name,
            deep_supervision=args.deep_supervision,
            deep_supr_num=args.deep_supr_num,
            res_block=args.res_block,
            trans_bias=True,)
        print(f"Filters: {model.filters},\nKernels: {kernels}\nStrides: {strides}")

    else:
        raise ValueError("Unsupported model " + str(args.model_name))

    model_dict = torch.load(pretrained_pth)
    model.load_state_dict(model_dict['state_dict'])
    model.eval()
    model.to(device)
    models.append(model)

# ---------------------------- ENSEMBLE EVALUATION ---------------------------- #

test_inferer = monai.inferers.SlidingWindowInferer(roi_size=(args.roi_x, args.roi_y, args.roi_z),
                                                   sw_batch_size=1,
                                                   overlap=args.infer_overlap)

metrics_saver = MetricsSaver(
    save_dir=output_directory,
    summary_ops=["mean", "median", "max", "5percentile", "95percentile", "notnans"],
)

mean_post_transforms = Compose(
    [
        EnsureTyped(keys=["pred0", "pred1", "pred2", "pred3", "pred4"]),
        MeanEnsembled(
            keys=["pred0", "pred1", "pred2", "pred3", "pred4"],
            output_key="pred",
            weights=[0.95, 0.94, 0.95, 0.94, 0.90],
        ),
        # Activationsd(keys="pred", softmax=True),
        # AsDiscreted(keys="pred", threshold=0.5),
        AsDiscreted(keys="pred", argmax=True, to_onehot=args.out_channels),
        AsDiscreted(keys="label", to_onehot=args.out_channels),  # manca anche decollate batch?
      #  SaveImaged(keys="pred", meta_keys="image", output_dir=output_directory, output_postfix="mean",
              #     output_ext=".nii.gz", resample=False)

    ]
)

logging.basicConfig(stream=sys.stdout, level=logging.INFO)
ensemble_evaluate(mean_post_transforms, models, torch.device(device), test_loader,
                  test_inferer, metrics_saver)

[KumoLiu]

Just to confirm, I meant batch_size in dataloader not sw_batch_size.

[casarinsof]

yes, also that is equal to 1. Thanks