Auto3DSeg: code for standalone training step adding algo filtering and skipping trained folds

[pwrightkcl]

Because of my local setup, I need to run the steps in AutoRunner separately, so I have written my code based on the examples in the MONAI Auto3DSeg Reference Python APIs notebook (section heading "Training the neural networks sequentially"). I noticed that the training step trains all algos and starts trained folds from scratch, so I added some lines to:

• specify algos to train
• skip trained folds

I'm new to this so I'm posting it here for discussion rather than doing a PR on the documentation. Comments, please!

import argparse
from pathlib import Path
from multiprocessing import freeze_support

from monai.apps.auto3dseg import (
    import_bundle_algo_history,
)
from monai.auto3dseg import algo_to_pickle
from monai.utils.enums import AlgoKeys


def main(work_dir: Path, algos: (list[str], None) = None, verbose: bool = False):
    freeze_support()

    print(f"Training models in {work_dir}.")

    if not work_dir.exists():
        raise FileNotFoundError(f"Could not find work directory {work_dir}")

    history = import_bundle_algo_history(str(work_dir), only_trained=False)
    print(f"Found {len(history)} model instances.")

    if algos:
        print(f"Filtering algos by name(s): {algos}.")
        history_in = [algo_dict for algo_dict in history if algo_dict[AlgoKeys.ALGO].name in algos]
        if len(history_in) == 0:
            raise ValueError(f"Could not find any algos with name(s): {algos}.")
        if verbose:
            algo_names_in = [algo_dict[AlgoKeys.ALGO].name for algo_dict in history_in]
            print("Instances found:")
            for algo in algos:
                print(f"{algo}: {algo_names_in.count(algo)}")
            history_out = [algo_dict for algo_dict in history if algo_dict[AlgoKeys.ALGO].name not in algos]
            if history_out:
                print("Excluded these instances not in algo list:")
                algo_names_out = [algo_dict[AlgoKeys.ALGO].name for algo_dict in history_out]
                for algo_out in set(algo_names_out):
                    print(f"{algo_out}: {algo_names_out.count(algo_out)}")
        history = history_in

    if any([algo_dict[AlgoKeys.IS_TRAINED] for algo_dict in history]):
        if verbose:
            trained_ids = [algo_dict[AlgoKeys.ID] for algo_dict in history if algo_dict[AlgoKeys.IS_TRAINED]]
            print(f"Skipping trained model instances: {trained_ids}.")
        history = [algo_dict for algo_dict in history if not algo_dict[AlgoKeys.IS_TRAINED]]

    train_param = {}
    # Add training parameter overrides here. Example:
    # train_param = {"learning_rate": 4.0e-5}

    for algo_dict in history:
        algo = algo_dict[AlgoKeys.ALGO]
        algo.train(train_param)
        acc = algo.get_score()
        algo_to_pickle(algo, template_path=algo.template_path, best_metric=acc)


if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument("work_dir", type=Path, help="Path to work directory (required)")
    parser.add_argument("--algos", nargs="*", type=str, help="Zero or more algo names to train (optional)")
    parser.add_argument("-v", "--verbose", action="store_true", help="Print verbose output (optional)")
    args = parser.parse_args()
    main(work_dir=args.work_dir, algos=args.algos, verbose=args.verbose)

[KumoLiu]

Hi @pwrightkcl, thanks for sharing, I think AutoRunner can help you:

1. train from specified algos.
   

   MONAI/monai/apps/auto3dseg/auto_runner.py

   Line 205 in f4103c5

   algos: dict | list | str | None = None,

2. check the history and training based on cache.
   

   MONAI/monai/apps/auto3dseg/auto_runner.py

   Line 847 in f4103c5

   history = import_bundle_algo_history(self.work_dir, only_trained=False)

Perhaps you're experiencing any issues with it? Or it still doesn't meet your needs? Thanks.

[pwrightkcl]

For my use case I need to run the AutoRunner steps separately, so this is to help people who have that constraint, since the AutoRunner features to select algos and skip trained folds aren't included in the tutorial with the separate steps.

[KumoLiu]

Hi @pwrightkcl, got it.
You can also use 'AutoRunner' to simplify your code. If this is the case, you can select the 'True' option for the step you need to do.

MONAI/monai/apps/auto3dseg/auto_runner.py

Lines 206 to 211 in f4103c5

	analyze: bool | None = None,
	algo_gen: bool | None = None,
	train: bool | None = None,
	hpo: bool = False,
	hpo_backend: str = "nni",
	ensemble: bool = True,

MONAI/monai/apps/auto3dseg/auto_runner.py

Line 795 in f4103c5

# step 1: data analysis

MONAI/monai/apps/auto3dseg/auto_runner.py

Line 811 in f4103c5

if self.algo_gen:

MONAI/monai/apps/auto3dseg/auto_runner.py

Line 845 in f4103c5

auto_train_choice = self.train is None

MONAI/monai/apps/auto3dseg/auto_runner.py

Line 875 in f4103c5

if self.ensemble:

Hope it helps, thanks.

[pwrightkcl]

Hi @KumoLiu . I had no idea you could just use kwargs to get AutoRunner to run specific steps. This pretty much makes my code redundant, unless you want silly levels of logging. Still, anyone finding this post in future should find the info useful.

It would be nice if these and other kwargs were in the documentation. I realise it's not possible to be exhaustive, but I've spend quite a while digging through code, tracing execution through all the various layers of class definitions, just to find the kwarg I need.

[KumoLiu]

Hi @pwrightkcl ,

It would be nice if these and other kwargs were in the documentation. I realise it's not possible to be exhaustive, but I've spend quite a while digging through code, tracing execution through all the various layers of class definitions, just to find the kwarg I need.

Your suggestion to improve the documentation regarding various keyword arguments is well-received and we understand it can be rather complex navigating through multiple layers of class definitions to find needed arguments.
We do indeed agree that our documentation can benefit from enhancements. If you wish to contribute a PR for this improvement, we would greatly appreciate it.
Additionally, you might find it of interest to know that we have another open ticket currently tracking parameters that can be manually overridden in Auto3dSeg.

Thank you again for your valuable feedback and suggestions.