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
