From 1fd8b2de5d0a5043cb173cf41d7146bb2d281573 Mon Sep 17 00:00:00 2001 From: Anwai Archit <52396323+anwai98@users.noreply.github.com> Date: Mon, 1 Apr 2024 22:47:34 +0200 Subject: [PATCH] Add scripts for training all benchmarking experiments (#17) --- .../vision-mamba/run_benchmarking_training.py | 109 +++++++++++++++++ .../unetr/run_benchmarking_training.py | 110 ++++++++++++++++++ 2 files changed, 219 insertions(+) create mode 100644 experiments/vision-mamba/run_benchmarking_training.py create mode 100644 experiments/vision-transformer/unetr/run_benchmarking_training.py diff --git a/experiments/vision-mamba/run_benchmarking_training.py b/experiments/vision-mamba/run_benchmarking_training.py new file mode 100644 index 00000000..4284f6b0 --- /dev/null +++ b/experiments/vision-mamba/run_benchmarking_training.py @@ -0,0 +1,109 @@ +import os +import shutil +import subprocess +from glob import glob +from datetime import datetime + + +SCRIPT_NAMES = { + "livecell": "livecell/run_livecell.py", + "cremi": "cremi/run_cremi.py", + "neurips-cellseg": "neurips-cellseg/run_neurips_cellseg.py", +} +ALL_MODELS = ["vim_t", "vim_s", "vim_b"] +ALL_SETTINGS = ["boundaries", "affinities", "distances"] + + +def write_batch_script(out_path, dataset, model_type, mode, setting): + batch_script = f"""#!/bin/bash +#SBATCH -t 2-00:00:00 +#SBATCH --nodes=1 +#SBATCH --ntasks=1 +#SBATCH -p grete:shared +#SBATCH -G A100:1 +#SBATCH -c 16 +#SBATCH -A gzz0001 +#SBATCH -x ggpu136 +#SBATCH --constraint=80gb +#SBATCH --job-name=vimunet-{dataset} + +source ~/.bashrc +mamba activate vm3 +""" + + python_script = f"python {SCRIPT_NAMES[dataset]} " + python_script += f"--{mode} " # train or predict + python_script += f"--{setting} " # boundaries / affinities / distances + python_script += f"-m {model_type} " + + batch_script += python_script + + _op = out_path[:-3] + f"{model_type}-{setting}-{dataset}.sh" + + with open(_op, "w") as f: + f.write(batch_script) + + +def write_all_scripts(batch_script, mode, dataset=None, model_type=None, setting=None): + if dataset is None: + dataset = list(SCRIPT_NAMES.keys()) + else: + dataset = [dataset] + + if model_type is None: + model_type = ALL_MODELS + else: + model_type = [model_type] + + if setting is None: + setting = ALL_SETTINGS + else: + setting = [setting] + + for _dataset in dataset: + for _model in model_type: + for _setting in setting: + write_batch_script( + out_path=batch_script, + dataset=_dataset, + model_type=_model, + mode=mode, + setting=_setting + ) + + +def submit_slurm(args): + """Submit python script that needs gpus with given inputs on a slurm node. + """ + tmp_folder = os.path.expanduser("./gpu_jobs") + os.makedirs(tmp_folder, exist_ok=True) + + script_name = "vimunet_" + dt = datetime.now().strftime("%Y_%m_%d_%H_%M_%S_%f") + tmp_name = script_name + dt + + batch_script = os.path.join(tmp_folder, f"{tmp_name}.sh") + + write_all_scripts( + batch_script, mode=args.phase, dataset=args.dataset, model_type=args.model_type, setting=args.setting + ) + + for my_script in glob(tmp_folder + "/*"): + cmd = ["sbatch", my_script] + subprocess.run(cmd) + + +if __name__ == "__main__": + try: + shutil.rmtree("./gpu_jobs") + except FileNotFoundError: + pass + + import argparse + parser = argparse.ArgumentParser() + parser.add_argument("-d", "--dataset", type=str, default=None) + parser.add_argument("-m", "--model_type", type=str, default=None) + parser.add_argument("-p", "--phase", type=str, default=None, required=True) + parser.add_argument("--setting", type=str, default=None) + args = parser.parse_args() + submit_slurm(args) diff --git a/experiments/vision-transformer/unetr/run_benchmarking_training.py b/experiments/vision-transformer/unetr/run_benchmarking_training.py new file mode 100644 index 00000000..022f4942 --- /dev/null +++ b/experiments/vision-transformer/unetr/run_benchmarking_training.py @@ -0,0 +1,110 @@ +import os +import shutil +import subprocess +from glob import glob +from datetime import datetime + + +SCRIPT_NAMES = { + "livecell": "livecell/run_livecell_for_vimunet.py", + "cremi": "cremi/run_cremi.py", + "neurips-cellseg": "neurips-cellseg/run_neurips_cellseg.py", +} +ALL_MODELS = ["unet", "vit_b", "vit_l", "vit_h"] +ALL_SETTINGS = ["boundaries", "affinities", "distances"] + + +def write_batch_script(out_path, dataset, model_type, mode, setting): + model_name = "unetr" if model_type.startswith("vit") else model_type + batch_script = f"""#!/bin/bash +#SBATCH -t 2-00:00:00 +#SBATCH --nodes=1 +#SBATCH --ntasks=1 +#SBATCH -p grete:shared +#SBATCH -G A100:1 +#SBATCH -c 16 +#SBATCH -A gzz0001 +#SBATCH -x ggpu136 +#SBATCH --constraint=80gb +#SBATCH --job-name={model_name}-{dataset} + +source ~/.bashrc +source activate sam +""" + + python_script = f"python {SCRIPT_NAMES[dataset]} " + python_script += f"--{mode} " # train or predict + python_script += f"--{setting} " # boundaries / affinities / distances + python_script += f"-m {model_type} " + + batch_script += python_script + + _op = out_path[:-3] + f"{model_type}-{setting}-{dataset}.sh" + + with open(_op, "w") as f: + f.write(batch_script) + + +def write_all_scripts(batch_script, mode, dataset=None, model_type=None, setting=None): + if dataset is None: + dataset = list(SCRIPT_NAMES.keys()) + else: + dataset = [dataset] + + if model_type is None: + model_type = ALL_MODELS + else: + model_type = [model_type] + + if setting is None: + setting = ALL_SETTINGS + else: + setting = [setting] + + for _dataset in dataset: + for _model in model_type: + for _setting in setting: + write_batch_script( + out_path=batch_script, + dataset=_dataset, + model_type=_model, + mode=mode, + setting=_setting + ) + + +def submit_slurm(args): + """Submit python script that needs gpus with given inputs on a slurm node. + """ + tmp_folder = os.path.expanduser("./gpu_jobs") + os.makedirs(tmp_folder, exist_ok=True) + + script_name = "unetr_" + dt = datetime.now().strftime("%Y_%m_%d_%H_%M_%S_%f") + tmp_name = script_name + dt + + batch_script = os.path.join(tmp_folder, f"{tmp_name}.sh") + + write_all_scripts( + batch_script, mode=args.phase, dataset=args.dataset, model_type=args.model_type, setting=args.setting + ) + + for my_script in glob(tmp_folder + "/*"): + cmd = ["sbatch", my_script] + subprocess.run(cmd) + + +if __name__ == "__main__": + try: + shutil.rmtree("./gpu_jobs") + except FileNotFoundError: + pass + + import argparse + parser = argparse.ArgumentParser() + parser.add_argument("-d", "--dataset", type=str, default=None) + parser.add_argument("-m", "--model_type", type=str, default=None) + parser.add_argument("-p", "--phase", type=str, default=None, required=True) + parser.add_argument("--setting", type=str, default=None) + args = parser.parse_args() + submit_slurm(args)