def train_eval_dealiasers(contrast='CORPD_FBK', n_epochs=200, n_samples=None, model_name=None, model_size=None, loss='mae'): job_name = 'dealiasing_fastmri' model_specs = list(get_model_specs(force_res=True, dealiasing=True)) if model_name is not None: model_specs = [ms for ms in model_specs if ms[0] == model_name] if model_size is not None: model_specs = [ms for ms in model_specs if ms[1] == model_size] n_models = len(model_specs) train_cluster = SLURMCluster( cores=1, job_cpu=20, memory='80GB', job_name=job_name, walltime='20:00:00', interface='ib0', job_extra=[ f'--gres=gpu:1', '--qos=qos_gpu-t3', '--distribution=block:block', '--hint=nomultithread', '--output=%x_%j.out', ], env_extra=[ 'cd $WORK/fastmri-reproducible-benchmark', '. ./submission_scripts_jean_zay/env_config.sh', ], ) train_cluster.adapt(minimum_jobs=0, maximum_jobs=n_models) client = Client(train_cluster) futures = [client.submit( # function to execute train_dealiaser, model_fun=model_fun, model_kwargs=kwargs, run_id=f'{model_name}_{model_size}', n_scales=n_scales, contrast=contrast, n_epochs=n_epochs, n_samples=n_samples, loss=loss, ) for model_name, model_size, model_fun, kwargs, _, n_scales, _ in model_specs] run_ids = client.gather(futures) client.close() train_cluster.close() # eval eval_dealiasers( run_ids, job_name=job_name, contrast=contrast, n_epochs=n_epochs, model_name=model_name, model_size=model_size, n_samples_train=n_samples, loss=loss, ) return run_ids
def train_xpdnet_click( model_name, model_size, af, brain, loss, refine_smaps, refine_big, n_epochs, checkpoint_epoch, n_epochs_original, save_state, original_run_id, contrast, equidistant_fake, n_iter, multi_gpu, ): n_primal = 5 model_fun, kwargs, n_scales, res = [ (model_fun, kwargs, n_scales, res) for m_name, m_size, model_fun, kwargs, _, n_scales, res in get_model_specs(n_primal=n_primal, force_res=False) if m_name == model_name and m_size == model_size ][0] train_xpdnet( model_fun=model_fun, model_kwargs=kwargs, model_size=model_size, multicoil=True, af=af, brain=brain, res=res, loss=loss, n_iter=n_iter, refine_smaps=refine_smaps or refine_big, refine_big=refine_big, n_scales=n_scales, n_primal=n_primal, n_epochs=n_epochs, checkpoint_epoch=checkpoint_epoch, n_epochs_original=n_epochs_original, save_state=save_state, original_run_id=original_run_id, contrast=contrast, equidistant_fake=equidistant_fake, multi_gpu=multi_gpu, )
def test_train_denoiser(create_full_fastmri_test_tmp_dataset): config_data.FASTMRI_DATA_DIR = create_full_fastmri_test_tmp_dataset['fastmri_tmp_data_dir'] config_data.PATHS_MAP[False][False]['train'] = 'singlecoil_train/singlecoil_train' config_output.LOGS_DIR = create_full_fastmri_test_tmp_dataset['logs_tmp_dir'] config_output.CHECKPOINTS_DIR = create_full_fastmri_test_tmp_dataset['checkpoints_tmp_dir'] from fastmri_recon.training_scripts.denoising.generic_train import train_denoiser model_specs = get_model_specs(force_res=True) model_specs = [ms for ms in model_specs if ms[1] == 'small'] for model_name, model_size, model_fun, kwargs, n_inputs, _, _ in model_specs: train_denoiser( model=(model_fun, kwargs, n_inputs), run_id=f'{model_name}_{model_size}_{int(time.time())}', n_epochs=1, n_steps_per_epoch=2, )
def test_train_dealiaser(create_full_fastmri_test_tmp_dataset): config.FASTMRI_DATA_DIR = create_full_fastmri_test_tmp_dataset[ 'fastmri_tmp_data_dir'] config.LOGS_DIR = create_full_fastmri_test_tmp_dataset['logs_tmp_dir'] config.CHECKPOINTS_DIR = create_full_fastmri_test_tmp_dataset[ 'checkpoints_tmp_dir'] from fastmri_recon.training_scripts.dealias_train import train_dealiaser model_specs = get_model_specs(force_res=True, dealiasing=True) model_specs = [ms for ms in model_specs if ms[1] == 'small'] for model_name, model_size, model_fun, kwargs, _, n_scales, _ in model_specs: train_dealiaser( model_fun=model_fun, model_kwargs=kwargs, n_scales=n_scales, run_id=f'{model_name}_{model_size}', n_epochs=1, n_steps_per_epoch=2, )
def evaluate_xpdnet_click( model_name, model_size, run_id, n_epochs, contrast, af, n_iter, cuda_visible_devices, n_samples, refine_smaps, refine_big, brain, equidistant_fake, ): n_primal = 5 model_fun, kwargs, n_scales, res = [ (model_fun, kwargs, n_scales, res) for m_name, m_size, model_fun, kwargs, _, n_scales, res in get_model_specs(n_primal=n_primal, force_res=False) if m_name == model_name and m_size == model_size ][0] metrics_names, eval_res = evaluate_xpdnet( model_fun=model_fun, model_kwargs=kwargs, multicoil=True, run_id=run_id, n_epochs=n_epochs, contrast=contrast, af=af, n_iter=n_iter, n_primal=n_primal, n_scales=n_scales, n_samples=n_samples, cuda_visible_devices=cuda_visible_devices, refine_smaps=refine_smaps or refine_big, refine_big=refine_big, brain=brain, equidistant_fake=equidistant_fake, ) print(metrics_names) print(eval_res)
def xpdnet_inference_click( model_name, model_size, af, n_iter, brain, challenge, refine_smaps, refine_big, n_epochs, run_id, exp_id, contrast, ): n_primal = 5 model_fun, kwargs, n_scales, res = [ (model_fun, kwargs, n_scales, res) for m_name, m_size, model_fun, kwargs, _, n_scales, res in get_model_specs(n_primal=n_primal, force_res=False) if m_name == model_name and m_size == model_size ][0] xpdnet_inference( model_fun=model_fun, model_kwargs=kwargs, af=af, n_iter=n_iter, brain=brain, challenge=challenge, refine_smaps=refine_smaps or refine_big, refine_big=refine_big, n_epochs=n_epochs, run_id=run_id, exp_id=exp_id, contrast=contrast, res=res, n_scales=n_scales, n_primal=n_primal, )
def train_eval_plug_and_play( contrast='CORPD_FBK', n_epochs=200, n_samples=None, af=4, n_primal=5, loss='compound_mssim', train_partition='gpu_p1', model_name=None, model_size=None, ): job_name = 'plug_and_play' model_specs = list(get_model_specs(force_res=False, n_primal=n_primal)) if model_name is not None: model_specs = [ms for ms in model_specs if ms[0] == model_name] if model_size is not None: model_specs = [ms for ms in model_specs if ms[1] == model_size] n_models = len(model_specs) train_cluster = SLURMCluster( cores=1, job_cpu=20, memory='80GB', job_name=job_name, walltime='60:00:00', interface='ib0', job_extra=[ '--gres=gpu:1', '--qos=qos_gpu-t4', '--distribution=block:block', '--hint=nomultithread', '--output=%x_%j.out', f'--partition {train_partition}', ], env_extra=[ 'cd $WORK/fastmri-reproducible-benchmark', '. ./submission_scripts_jean_zay/env_config.sh', ], ) train_cluster.adapt(minimum_jobs=0, maximum_jobs=n_models) client = Client(train_cluster) futures = [ client.submit( # function to execute train_xpdnet, model_fun=model_fun, model_kwargs=kwargs, model_size=model_size, multicoil=False, n_scales=n_scales, res=res, n_primal=n_primal, contrast=contrast, n_epochs=n_epochs, n_samples=n_samples, af=af, loss=loss, ) for _, model_size, model_fun, kwargs, _, n_scales, res in model_specs ] run_ids = client.gather(futures) client.close() train_cluster.close() # eval eval_plug_and_play( run_ids, job_name=job_name, contrast=contrast, n_epochs=n_epochs, af=af, n_primal=n_primal, model_name=model_name, model_size=model_size, n_samples_train=n_samples, ) return run_ids
def eval_plug_and_play( run_ids, job_name='eval_pandp', contrast='CORPD_FBK', n_samples_train=None, n_epochs=200, af=4, n_primal=5, train_partition='gpu_p1', model_name=None, model_size=None, ): model_specs = list(get_model_specs(force_res=False, n_primal=n_primal)) if model_name is not None: model_specs = [ms for ms in model_specs if ms[0] == model_name] if model_size is not None: model_specs = [ms for ms in model_specs if ms[1] == model_size] n_models = len(model_specs) # eval eval_cluster = SLURMCluster( cores=1, job_cpu=40, memory='80GB', job_name=job_name, walltime='2:00:00', interface='ib0', job_extra=[ f'--gres=gpu:1', '--qos=qos_gpu-t3', '--distribution=block:block', '--hint=nomultithread', '--output=%x_%j.out', ], env_extra=[ 'cd $WORK/fastmri-reproducible-benchmark', '. ./submission_scripts_jean_zay/env_config.sh', ], ) eval_cluster.adapt(minimum_jobs=0, maximum_jobs=n_models) client = Client(eval_cluster) futures = [ client.submit( # function to execute evaluate_xpdnet, model_fun=model_fun, model_kwargs=kwargs, run_id=run_id, multicoil=False, n_samples=50, contrast=contrast, af=af, n_epochs=n_epochs, n_scales=n_scales, res=res, ) for run_id, (_, _, model_fun, kwargs, _, n_scales, res) in zip(run_ids, model_specs) ] df_results = pd.DataFrame( columns='model_name model_size psnr ssim'.split()) for (name, model_size, _, _, _, _, _), future in zip(model_specs, futures): _, eval_res = client.gather(future) df_results = df_results.append(dict( model_name=name, model_size=model_size, psnr=eval_res[0], ssim=eval_res[1], ), ignore_index=True) print(df_results) outputs_file = f'reconstruction_results_{n_samples_train}.csv' if model_name is not None: outputs_file = f'reconstruction_results_{n_samples_train}_{model_name}.csv' df_results.to_csv(outputs_file) print('Shutting down dask workers') client.close() eval_cluster.close() return run_ids
from jean_zay.submitit.fastmri_reproducible_benchmark.mem_fitting_test import test_works_in_xpdnet_train n_iter_to_try_for_size = { 'medium': range(30, 40), } n_primal = 5 job_name = 'xpdnet_tryouts' executor = get_executor(job_name, timeout_hour=1, n_gpus=1, project='fastmri4') results = {} with executor.batch(): for data_consistency_learning in [True, False]: for model_size_spec, n_iter_to_try in n_iter_to_try_for_size.items(): for model_name, model_size, model_fun, model_kwargs, n_inputs, n_scales, res in get_model_specs( n_primal): if model_size != model_size_spec or model_name != 'MWCNN': continue for n_iter in n_iter_to_try: job = executor.submit( test_works_in_xpdnet_train, model_fun=model_fun, model_kwargs=model_kwargs, n_scales=n_scales, res=res, n_iter=n_iter, multicoil=True, use_mixed_precision=True, data_consistency_learning=data_consistency_learning, ) res_id = (model_name, model_size,
loss = 'compound_mssim' lr = 1e-4 batch_size = None n_samples = None n_epochs = 500 n_primal = 5 contrast = None refine_smaps = True refine_big = True n_dual = 2 primal_only = False multiscale_kspace_learning = False n_dual_filters = 8 n_iter = 18 use_mixed_precision = False model_specs = list(get_model_specs(force_res=False, n_primal=n_primal)) if model_name is not None: model_specs = [ms for ms in model_specs if ms[0] == model_name] if model_size is not None: model_specs = [ms for ms in model_specs if ms[1] == model_size] parameter_grid = [ dict( model_fun=model_fun, model_kwargs=kwargs, model_size=model_size, multicoil=True, n_scales=n_scales, res=res, n_primal=n_primal, contrast=contrast,
from fastmri_recon.evaluate.scripts.denoising_eval import evaluate_xpdnet_denoising from fastmri_recon.models.subclassed_models.denoisers.proposed_params import get_model_specs from fastmri_recon.training_scripts.denoising.generic_train import train_denoiser import pandas as pd from jean_zay.submitit.general_submissions import train_eval_grid, eval_grid job_name = 'denoising_fastmri' model_name = None model_size = None loss = 'compound_mssim' n_samples = None n_epochs = 800 contrast = 'CORPD_FBK' model_specs = list(get_model_specs(force_res=True)) if model_name is not None: model_specs = [ms for ms in model_specs if ms[0] == model_name] if model_size is not None: model_specs = [ms for ms in model_specs if ms[1] == model_size] parameter_grid = [ dict( model=(model_fun, kwargs, n_inputs), run_id=f'{model_name}_{model_size}_{int(time.time())}', contrast=contrast, n_epochs=n_epochs, n_samples=n_samples, loss=loss, noise_std=2, ) for model_name, model_size, model_fun, kwargs, n_inputs, _, _ in