def _load_net(path, subnet, subnet_params, it_net_params):
subnet = subnet(**subnet_params).to(device)
it_net = IterativeNet(subnet, **it_net_params).to(device)
it_net.load_state_dict(torch.load(path, map_location=torch.device(device)))
it_net.freeze()
it_net.eval()
return it_net
# ------ construct network and train -----
subnet_tmp = subnet(**subnet_params).to(device)
it_net_tmp = IterativeNet(
subnet_tmp, **{
"num_iter": 8,
"lam": 8 * [0.1],
"lam_learnable": False,
"final_dc": False,
"resnet_factor": 1.0,
"concat_mask": False,
"multi_slice": False,
}).to(device)
it_net_tmp.load_state_dict(
torch.load(
"results/radial_50_no_fs_unet_it_preinit_v1_train_phase_1/" +
"model_weights.pt",
map_location=torch.device(device),
))
subnet = it_net_tmp.subnet
it_net = IterativeNet(subnet, **it_net_params).to(device)
train_data = train_data("train", **train_data_params)
val_data = val_data("val", **val_data_params)
for i in range(train_phases):
train_params_cur = {}
for key, value in train_params.items():
train_params_cur[key] = (value[i] if isinstance(value,
(tuple,
list)) else value)
# ------ construct network and train -----
subnet_tmp = subnet(**subnet_params).to(device)
it_net_tmp = IterativeNet(
subnet_tmp, **{
"num_iter": 1,
"lam": 0.0,
"lam_learnable": False,
"final_dc": False,
"resnet_factor": 1.0,
"concat_mask": False,
"multi_slice": False,
"ee": mask,
}).to(device)
it_net_tmp.load_state_dict(
torch.load(
"results/radial_50_no_fs_tiramisu_ee_v4_train_phase_1/" +
"model_weights.pt",
map_location=torch.device(device),
))
subnet = it_net_tmp.subnet
it_net = IterativeNet(subnet, **it_net_params).to(device)
train_data = train_data("train", **train_data_params)
val_data = val_data("val", **val_data_params)
for i in range(train_phases):
train_params_cur = {}
for key, value in train_params.items():
train_params_cur[key] = (value[i] if isinstance(value,
(tuple,
list)) else value)
subnet_tmp = subnet(**subnet_params).to(device)
it_net_tmp = IterativeNet(
subnet_tmp,
**{
"num_iter": 1,
"lam": 0.0,
"lam_learnable": False,
"final_dc": False,
"resnet_factor": 1.0,
"operator": OpA,
"inverter": inverter,
}
).to(device)
it_net_tmp.load_state_dict(
torch.load(
"results/Fourier_UNet_jitter_v3_train_phase_2/model_weights.pt",
map_location=torch.device(device),
)
)
subnet = it_net_tmp.subnet
it_net = IterativeNet(subnet, **it_net_params).to(device)
train_data = train_data("train", **train_data_params)
val_data = val_data("val", **val_data_params)
for i in range(train_phases):
train_params_cur = {}
for key, value in train_params.items():
train_params_cur[key] = (
value[i] if isinstance(value, (tuple, list)) else value
)
# ------ construct network and train -----
subnet_tmp = subnet(**subnet_params).to(device)
it_net_tmp = IterativeNet(
subnet_tmp, **{
"num_iter": 1,
"lam": 0.0,
"lam_learnable": False,
"final_dc": False,
"resnet_factor": 1.0,
"operator": OpA,
"inverter": OpA.inv,
}).to(device)
it_net_tmp.load_state_dict(
torch.load(
"results/Radon_Tiramisu_jitter_v4_train_phase_1/model_weights.pt",
map_location=torch.device(device),
))
subnet = it_net_tmp.subnet
it_net = IterativeNet(subnet, **it_net_params).to(device)
train_data = train_data("train", **train_data_params)
val_data = val_data("val", **val_data_params)
for i in range(train_phases):
train_params_cur = {}
for key, value in train_params.items():
train_params_cur[key] = (value[i] if isinstance(value,
(tuple,
list)) else value)
for key, value in it_net_params.items():
file.write(key + ": " + str(value) + "\n")
for key, value in train_params.items():
file.write(key + ": " + str(value) + "\n")
for key, value in train_data_params.items():
file.write(key + ": " + str(value) + "\n")
for key, value in val_data_params.items():
file.write(key + ": " + str(value) + "\n")
file.write("train_phases" + ": " + str(train_phases) + "\n")
# ------ construct network and train -----
subnet = subnet(**subnet_params).to(device)
it_net = IterativeNet(subnet, **it_net_params).to(device)
it_net.load_state_dict(
torch.load(
"results/Fourier_UNet_it_jit-nojit_pre_train_phase_1/model_weights.pt",
map_location=torch.device(device),
)
)
train_data = train_data("train", **train_data_params)
val_data = val_data("val", **val_data_params)
for i in range(train_phases):
train_params_cur = {}
for key, value in train_params.items():
train_params_cur[key] = (
value[i] if isinstance(value, (tuple, list)) else value
)
print("Phase {}:".format(i + 1))
for key, value in train_params_cur.items():
# ------ construct network and train -----
subnet_tmp = subnet(**subnet_params).to(device)
it_net_tmp = IterativeNet(
subnet_tmp, **{
"num_iter": 1,
"lam": 0.0,
"lam_learnable": False,
"final_dc": False,
"resnet_factor": 1.0,
"operator": OpA,
"inverter": inverter,
}).to(device)
it_net_tmp.load_state_dict(
torch.load(
"results/Fourier_Tiramisu_jitter_v3_train_phase_1/" +
"model_weights_epoch40.pt",
map_location=torch.device(device),
))
subnet = it_net_tmp.subnet
it_net = IterativeNet(subnet, **it_net_params).to(device)
train_data = train_data("train", **train_data_params)
val_data = val_data("val", **val_data_params)
for i in range(train_phases):
train_params_cur = {}
for key, value in train_params.items():
train_params_cur[key] = (value[i] if isinstance(value,
(tuple,
list)) else value)