def __init__(
self,
subnet,
num_iter,
lam,
lam_learnable,
final_dc=True,
resnet_factor=1.0,
concat_mask=False,
multi_slice=False,
ee=False,
):
super(IterativeNet, self).__init__()
self.subnet = subnet
self.num_iter = num_iter
self.final_dc = final_dc
self.resnet_factor = resnet_factor
self.concat_mask = concat_mask
self.multi_slice = multi_slice
if not isinstance(lam, (list, tuple)):
lam = [lam] * num_iter
if not isinstance(lam_learnable, (list, tuple)):
lam_learnable = [lam_learnable] * len(lam)
self.lam = torch.nn.ParameterList([
torch.nn.Parameter(torch.tensor(lam[it]),
requires_grad=lam_learnable[it])
for it in range(len(lam))
])
# for fully learned nets
if isinstance(ee, torch.Tensor):
self.inverter = LearnableInverter(ee.shape[-2:],
ee,
learnable=True)
self.ee = ee
unprep_fft_channel,
)
# ----- load configuration -----
import config # isort:skip
# ----- global configuration -----
mpl.use("agg")
device = torch.device("cuda:0")
torch.cuda.set_device(0)
# ----- measurement configuration -----
mask_func = RadialMaskFunc(config.n, 40)
mask = unprep_fft_channel(mask_func((1, 1) + config.n + (1, )))
OpA = Fourier(mask)
inverter = LearnableInverter(config.n, mask, learnable=False)
# ----- network configuration -----
subnet_params = {
"in_channels": 2,
"out_channels": 2,
"drop_factor": 0.0,
"down_blocks": (5, 7, 9, 12, 15),
"up_blocks": (15, 12, 9, 7, 5),
"pool_factors": (2, 2, 2, 2, 2),
"bottleneck_layers": 20,
"growth_rate": 16,
"out_chans_first_conv": 16,
}
subnet = Tiramisu
"plt_marker": "o",
"plt_linestyle": ":",
"plt_linewidth": 2.75,
},
_load_net(
"results/Fourier_UNet_jitter_v3_train_phase_2/model_weights.pt",
UNet,
unet_params,
{
"num_iter": 1,
"lam": 0.0,
"lam_learnable": False,
"final_dc": False,
"resnet_factor": 1.0,
"operator": OpA,
"inverter": LearnableInverter(config.n, mask, learnable=False),
},
),
)
_append_net(
"UNet EE jit",
{
"name_disp": "UNetFL",
"name_save": "unet_ee_jit",
"plt_color": "maroon",
"plt_marker": "o",
"plt_linestyle": "-",
"plt_linewidth": None,
},
_load_net(
unprep_fft_channel,
)
# ----- load configuration -----
import config # isort:skip
# ----- global configuration -----
mpl.use("agg")
device = torch.device("cuda:0")
torch.cuda.set_device(0)
# ----- measurement configuration -----
mask_func = RadialMaskFunc(config.n, 40)
mask = unprep_fft_channel(mask_func((1, 1) + config.n + (1, )))
OpA = Fourier(mask)
inverter = LearnableInverter(config.n, mask, learnable=True)
# ----- network configuration -----
subnet_params = {
"in_channels": 2,
"drop_factor": 0.0,
"base_features": 32,
"out_channels": 2,
}
subnet = UNet
it_net_params = {
"num_iter": 1,
"lam": 0.0,
"lam_learnable": False,
"final_dc": False,