Demo Colibri.

In this example we show how to use a simple pipeline of end-to-end learning with the CASSI and SPC forward models. Mainly, the forward model is defined,

\[\mathbf{y} = \mathbf{H}_\phi \mathbf{x}\]

where \(\mathbf{H}\) is the forward model, \(\mathbf{x}\) is the input image and \(\mathbf{y}\) is the measurement and \(\phi\) are the coding elements of the forward model. The recovery model is defined as,

\[\mathbf{x} = \mathcal{G}_\theta( \mathbf{y})\]

where \(\mathcal{G}\) is the recovery model and \(\theta\) are the parameters of the recovery model.

The training is performed by minimizing the following loss function,

\[\{\phi^*,\theta^*\} = \arg \min_{\phi,\theta} \sum_{p=1}^{P}\mathcal{L}(\mathbf{x}_p, \mathcal{G}_\theta( \mathbf{H}_\phi \mathbf{x}_p)) + \lambda \mathcal{R}(\phi) + \mu \mathcal{R}(\mathbf{H}_\phi \mathbf{x})\]

where \(\mathcal{L}\) is the loss function, \(\mathcal{R}\) is the regularizer, \(\lambda\) and \(\mu\) are the regularization weights, and \(P\) is the number of samples in the training dataset.

Select Working Directory and Device

import os

os.chdir(os.path.dirname(os.getcwd()))
print("Current Working Directory ", os.getcwd())

# General imports
import matplotlib.pyplot as plt
import torch
from torch.utils.data import DataLoader

manual_device = "cpu"
# Check GPU support
print("GPU support: ", torch.cuda.is_available())

if manual_device:
    device = manual_device
else:
    device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
Current Working Directory  /home/runner/work/pycolibri/pycolibri
GPU support:  False

Load dataset

from colibri.data.datasets import CustomDataset

name = 'cifar10'  # ['cifar10', 'cifar100', 'mnist', 'fashion_mnist', 'cave']
path = '.'
batch_size = 128
acquisition_name = 'c_cassi'  # ['spc', 'cassi', 'doe']


dataset = CustomDataset(name, path)


dataset_loader = DataLoader(dataset, batch_size=batch_size, shuffle=False, num_workers=0)

Visualize dataset

from torchvision.utils import make_grid

sample = next(iter(dataset_loader))['input']
img = make_grid(sample[:32], nrow=8, padding=1, normalize=True, scale_each=False, pad_value=0)

plt.figure(figsize=(10, 10))
plt.imshow(img.permute(1, 2, 0))
plt.title('CIFAR10 dataset')
plt.axis('off')
plt.show()
CIFAR10 dataset

Optics forward model

Define the forward operators \(\mathbf{y} = \mathbf{H}_\phi \mathbf{x}\), in this case, the CASSI and SPC forward models. Each optics model can comptute the forward and backward operators i.e., \(\mathbf{y} = \mathbf{H}_\phi \mathbf{x}\) and \(\mathbf{x} = \mathbf{H}^T_\phi \mathbf{y}\).

import math
from colibri.optics import SPC, SD_CASSI, DD_CASSI, C_CASSI, SingleDOESpectral
from colibri.optics.sota_does import spiral_doe, spiral_refractive_index

img_size = sample.shape[1:]

acquisition_config = dict(
    input_shape=img_size,
)

if acquisition_name == 'spc':
    n_measurements  = 256
    n_measurements_sqrt = int(math.sqrt(n_measurements))
    acquisition_config['n_measurements'] = n_measurements

elif acquisition_name == 'doe':
    wavelengths = torch.Tensor([450, 550, 650])*1e-9
    doe_size = [100, 100]
    radius_doe = 0.5e-3
    source_distance = 1  # meters
    sensor_distance = 50e-3
    pixel_size = (2 * radius_doe) / min(doe_size)
    height_map, aperture = spiral_doe(M=doe_size[0], N=doe_size[1],
                                      number_spirals=3, radius=radius_doe,
                                      focal=50e-3, start_w=450e-9, end_w=650e-9)
    refractive_index = spiral_refractive_index

    acquisition_config.update({"height_map": height_map,
                               "aperture": aperture,
                               "wavelengths": wavelengths,
                               "source_distance": source_distance,
                               "sensor_distance": sensor_distance,
                               "sensor_spectral_sensitivity": lambda x: x,
                               "pixel_size": pixel_size,
                               "doe_refractive_index": refractive_index,
                               "trainable": True})

acquisition_model = {
    'spc': SPC,
    'sd_cassi': SD_CASSI,
    'dd_cassi': DD_CASSI,
    'c_cassi': C_CASSI,
    'doe': SingleDOESpectral,
}[acquisition_name]

acquisition_model = acquisition_model(**acquisition_config)

y = acquisition_model(sample)

if acquisition_name == 'spc':
    y = y.reshape(y.shape[0], -1, n_measurements_sqrt, n_measurements_sqrt)

img = make_grid(y[:32], nrow=8, padding=1, normalize=True, scale_each=False, pad_value=0)

plt.figure(figsize=(10, 10))
plt.imshow(img.permute(1, 2, 0))
plt.axis('off')
plt.title(f'{acquisition_name.upper()} measurements')
plt.show()
C_CASSI measurements

Reconstruction model

Define the recovery model \(\mathbf{x} = \mathcal{G}_\theta( \mathbf{y})\), in this case, a simple U-Net model. You can add you custom model by using the :meth: build_network function. Additionally we define the end-to-end model that combines the forward and recovery models. Define the loss function \(\mathcal{L}\), and the regularizers \(\mathcal{R}\) for the forward and recovery models.

from colibri.models import build_network, Unet, Autoencoder
from colibri.misc import E2E
from colibri.train import Training
from colibri.metrics import psnr, ssim

from colibri.regularizers import (
    Binary,
    Transmittance,
    MinVariance,
    KLGaussian,
)

network_config = dict(
    in_channels=sample.shape[1],
    out_channels=sample.shape[1],
    reduce_spatial=True  # Only for Autoencoder
)

recovery_model = build_network(Unet, **network_config)

model = E2E(acquisition_model, recovery_model)
model = model.to(device)

optimizer = torch.optim.Adam(model.parameters(), lr=1e-3)
losses = {"MSE": torch.nn.MSELoss(), "L1": torch.nn.L1Loss()}
metrics = {"PSNR": psnr, "SSIM": ssim}
losses_weights = [1.0, 1.0]

n_epochs = 10
steps_per_epoch = 10
frequency = 1

if "cassi" in acquisition_name or "spc" in acquisition_name:
    regularizers_optics_ce = {"RB": Binary(), "RT": Transmittance()}
    regularizers_optics_ce_weights = [50, 1]
else:
    regularizers_optics_ce = {}
    regularizers_optics_ce_weights = []

regularizers_optics_mo = {"MV": MinVariance(), "KLG": KLGaussian(stddev=0.1)}
regularizers_optics_mo_weights = [1e-3, 0.1]

train_schedule = Training(
    model=model,
    train_loader=dataset_loader,
    optimizer=optimizer,
    loss_func=losses,
    losses_weights=losses_weights,
    metrics=metrics,
    regularizers=None,
    regularization_weights=None,
    schedulers=[],
    callbacks=[],
    device=device,
    regularizers_optics_ce=regularizers_optics_ce,
    regularization_optics_weights_ce=regularizers_optics_ce_weights,
    regularizers_optics_mo=regularizers_optics_mo,
    regularization_optics_weights_mo=regularizers_optics_mo_weights,
)

results = train_schedule.fit(
    n_epochs=n_epochs, steps_per_epoch=steps_per_epoch, freq=frequency
)
  0%|          | 0/391 [00:00<?, ?it/s]
Train :: Epoch: 1/10:   0%|          | 0/391 [00:00<?, ?it/s]
Train :: Epoch: 1/10:   0%|          | 1/391 [00:00<06:15,  1.04it/s, s=MSE: 0.08, L1: 0.22, , RB: 1725.91, RT: 6.24, , MV: 0.01, KLG: 4.63, , PSNR: 10.71, SSIM: 0.05,]
Train :: Epoch: 1/10:   1%|          | 2/391 [00:00<03:07,  2.07it/s, s=MSE: 0.08, L1: 0.22, , RB: 1725.91, RT: 6.24, , MV: 0.01, KLG: 4.63, , PSNR: 10.71, SSIM: 0.05,]
Train :: Epoch: 1/10:   1%|          | 2/391 [00:01<03:07,  2.07it/s, s=MSE: 0.05, L1: 0.18, , RB: 1725.91, RT: 6.24, , MV: 0.01, KLG: 4.67, , PSNR: 12.26, SSIM: 0.14,]
Train :: Epoch: 1/10:   1%|          | 3/391 [00:01<04:14,  1.52it/s, s=MSE: 0.05, L1: 0.18, , RB: 1725.91, RT: 6.24, , MV: 0.01, KLG: 4.67, , PSNR: 12.26, SSIM: 0.14,]
Train :: Epoch: 1/10:   1%|          | 3/391 [00:02<04:14,  1.52it/s, s=MSE: 0.04, L1: 0.16, , RB: 1725.91, RT: 6.24, , MV: 0.01, KLG: 4.81, , PSNR: 12.45, SSIM: 0.21,]
Train :: Epoch: 1/10:   1%|          | 4/391 [00:02<04:49,  1.34it/s, s=MSE: 0.04, L1: 0.16, , RB: 1725.91, RT: 6.24, , MV: 0.01, KLG: 4.81, , PSNR: 12.45, SSIM: 0.21,]
Train :: Epoch: 1/10:   1%|          | 4/391 [00:03<04:49,  1.34it/s, s=MSE: 0.04, L1: 0.16, , RB: 1725.91, RT: 6.24, , MV: 0.01, KLG: 4.72, , PSNR: 12.97, SSIM: 0.25,]
Train :: Epoch: 1/10:   1%|▏         | 5/391 [00:03<05:09,  1.25it/s, s=MSE: 0.04, L1: 0.16, , RB: 1725.91, RT: 6.24, , MV: 0.01, KLG: 4.72, , PSNR: 12.97, SSIM: 0.25,]
Train :: Epoch: 1/10:   1%|▏         | 5/391 [00:04<05:09,  1.25it/s, s=MSE: 0.03, L1: 0.14, , RB: 1725.91, RT: 6.24, , MV: 0.01, KLG: 4.71, , PSNR: 14.16, SSIM: 0.28,]
Train :: Epoch: 1/10:   2%|▏         | 6/391 [00:04<05:21,  1.20it/s, s=MSE: 0.03, L1: 0.14, , RB: 1725.91, RT: 6.24, , MV: 0.01, KLG: 4.71, , PSNR: 14.16, SSIM: 0.28,]
Train :: Epoch: 1/10:   2%|▏         | 6/391 [00:05<05:21,  1.20it/s, s=MSE: 0.03, L1: 0.13, , RB: 1725.91, RT: 6.24, , MV: 0.01, KLG: 4.88, , PSNR: 14.89, SSIM: 0.31,]
Train :: Epoch: 1/10:   2%|▏         | 7/391 [00:05<05:29,  1.17it/s, s=MSE: 0.03, L1: 0.13, , RB: 1725.91, RT: 6.24, , MV: 0.01, KLG: 4.88, , PSNR: 14.89, SSIM: 0.31,]
Train :: Epoch: 1/10:   2%|▏         | 7/391 [00:06<05:29,  1.17it/s, s=MSE: 0.03, L1: 0.13, , RB: 1725.91, RT: 6.24, , MV: 0.01, KLG: 4.83, , PSNR: 15.01, SSIM: 0.33,]
Train :: Epoch: 1/10:   2%|▏         | 8/391 [00:06<05:35,  1.14it/s, s=MSE: 0.03, L1: 0.13, , RB: 1725.91, RT: 6.24, , MV: 0.01, KLG: 4.83, , PSNR: 15.01, SSIM: 0.33,]
Train :: Epoch: 1/10:   2%|▏         | 8/391 [00:07<05:35,  1.14it/s, s=MSE: 0.02, L1: 0.12, , RB: 1725.91, RT: 6.24, , MV: 0.01, KLG: 4.67, , PSNR: 15.86, SSIM: 0.35,]
Train :: Epoch: 1/10:   2%|▏         | 9/391 [00:07<05:37,  1.13it/s, s=MSE: 0.02, L1: 0.12, , RB: 1725.91, RT: 6.24, , MV: 0.01, KLG: 4.67, , PSNR: 15.86, SSIM: 0.35,]
Train :: Epoch: 1/10:   2%|▏         | 9/391 [00:08<05:37,  1.13it/s, s=MSE: 0.02, L1: 0.12, , RB: 1725.91, RT: 6.24, , MV: 0.01, KLG: 4.86, , PSNR: 15.95, SSIM: 0.38,]
Train :: Epoch: 1/10:   3%|▎         | 10/391 [00:08<05:39,  1.12it/s, s=MSE: 0.02, L1: 0.12, , RB: 1725.91, RT: 6.24, , MV: 0.01, KLG: 4.86, , PSNR: 15.95, SSIM: 0.38,]
Train :: Epoch: 1/10:   3%|▎         | 10/391 [00:09<05:39,  1.12it/s, s=MSE: 0.02, L1: 0.12, , RB: 1725.91, RT: 6.24, , MV: 0.01, KLG: 4.85, , PSNR: 16.12, SSIM: 0.38,]
Train :: Epoch: 1/10:   3%|▎         | 11/391 [00:09<05:40,  1.12it/s, s=MSE: 0.02, L1: 0.12, , RB: 1725.91, RT: 6.24, , MV: 0.01, KLG: 4.85, , PSNR: 16.12, SSIM: 0.38,]
Train :: Epoch: 1/10:   3%|▎         | 11/391 [00:10<05:40,  1.12it/s, s=MSE: 0.02, L1: 0.11, , RB: 1725.91, RT: 6.24, , MV: 0.01, KLG: 4.69, , PSNR: 16.68, SSIM: 0.41,]
Train :: Epoch: 1/10:   3%|▎         | 11/391 [00:10<05:46,  1.10it/s, s=MSE: 0.02, L1: 0.11, , RB: 1725.91, RT: 6.24, , MV: 0.01, KLG: 4.69, , PSNR: 16.68, SSIM: 0.41,]

  0%|          | 0/391 [00:00<?, ?it/s]
Train :: Epoch: 2/10:   0%|          | 0/391 [00:00<?, ?it/s]
Train :: Epoch: 2/10:   0%|          | 1/391 [00:00<05:54,  1.10it/s, s=MSE: 0.02, L1: 0.11, , RB: 1725.91, RT: 6.24, , MV: 0.01, KLG: 4.63, , PSNR: 16.75, SSIM: 0.42,]
Train :: Epoch: 2/10:   1%|          | 2/391 [00:00<02:57,  2.20it/s, s=MSE: 0.02, L1: 0.11, , RB: 1725.91, RT: 6.24, , MV: 0.01, KLG: 4.63, , PSNR: 16.75, SSIM: 0.42,]
Train :: Epoch: 2/10:   1%|          | 2/391 [00:01<02:57,  2.20it/s, s=MSE: 0.02, L1: 0.10, , RB: 1725.91, RT: 6.24, , MV: 0.01, KLG: 4.67, , PSNR: 17.10, SSIM: 0.43,]
Train :: Epoch: 2/10:   1%|          | 3/391 [00:01<04:09,  1.56it/s, s=MSE: 0.02, L1: 0.10, , RB: 1725.91, RT: 6.24, , MV: 0.01, KLG: 4.67, , PSNR: 17.10, SSIM: 0.43,]
Train :: Epoch: 2/10:   1%|          | 3/391 [00:02<04:09,  1.56it/s, s=MSE: 0.02, L1: 0.11, , RB: 1725.91, RT: 6.24, , MV: 0.01, KLG: 4.81, , PSNR: 16.72, SSIM: 0.44,]
Train :: Epoch: 2/10:   1%|          | 4/391 [00:02<04:47,  1.35it/s, s=MSE: 0.02, L1: 0.11, , RB: 1725.91, RT: 6.24, , MV: 0.01, KLG: 4.81, , PSNR: 16.72, SSIM: 0.44,]
Train :: Epoch: 2/10:   1%|          | 4/391 [00:03<04:47,  1.35it/s, s=MSE: 0.02, L1: 0.11, , RB: 1725.91, RT: 6.24, , MV: 0.01, KLG: 4.72, , PSNR: 16.65, SSIM: 0.45,]
Train :: Epoch: 2/10:   1%|▏         | 5/391 [00:03<05:08,  1.25it/s, s=MSE: 0.02, L1: 0.11, , RB: 1725.91, RT: 6.24, , MV: 0.01, KLG: 4.72, , PSNR: 16.65, SSIM: 0.45,]
Train :: Epoch: 2/10:   1%|▏         | 5/391 [00:04<05:08,  1.25it/s, s=MSE: 0.02, L1: 0.10, , RB: 1725.91, RT: 6.24, , MV: 0.01, KLG: 4.71, , PSNR: 17.18, SSIM: 0.46,]
Train :: Epoch: 2/10:   2%|▏         | 6/391 [00:04<05:21,  1.20it/s, s=MSE: 0.02, L1: 0.10, , RB: 1725.91, RT: 6.24, , MV: 0.01, KLG: 4.71, , PSNR: 17.18, SSIM: 0.46,]
Train :: Epoch: 2/10:   2%|▏         | 6/391 [00:05<05:21,  1.20it/s, s=MSE: 0.02, L1: 0.10, , RB: 1725.91, RT: 6.24, , MV: 0.01, KLG: 4.88, , PSNR: 17.26, SSIM: 0.47,]
Train :: Epoch: 2/10:   2%|▏         | 7/391 [00:05<05:29,  1.17it/s, s=MSE: 0.02, L1: 0.10, , RB: 1725.91, RT: 6.24, , MV: 0.01, KLG: 4.88, , PSNR: 17.26, SSIM: 0.47,]
Train :: Epoch: 2/10:   2%|▏         | 7/391 [00:06<05:29,  1.17it/s, s=MSE: 0.02, L1: 0.10, , RB: 1725.91, RT: 6.24, , MV: 0.01, KLG: 4.83, , PSNR: 16.93, SSIM: 0.47,]
Train :: Epoch: 2/10:   2%|▏         | 8/391 [00:06<05:34,  1.15it/s, s=MSE: 0.02, L1: 0.10, , RB: 1725.91, RT: 6.24, , MV: 0.01, KLG: 4.83, , PSNR: 16.93, SSIM: 0.47,]
Train :: Epoch: 2/10:   2%|▏         | 8/391 [00:07<05:34,  1.15it/s, s=MSE: 0.02, L1: 0.09, , RB: 1725.91, RT: 6.24, , MV: 0.01, KLG: 4.67, , PSNR: 17.68, SSIM: 0.48,]
Train :: Epoch: 2/10:   2%|▏         | 9/391 [00:07<05:36,  1.14it/s, s=MSE: 0.02, L1: 0.09, , RB: 1725.91, RT: 6.24, , MV: 0.01, KLG: 4.67, , PSNR: 17.68, SSIM: 0.48,]
Train :: Epoch: 2/10:   2%|▏         | 9/391 [00:08<05:36,  1.14it/s, s=MSE: 0.02, L1: 0.09, , RB: 1725.91, RT: 6.24, , MV: 0.01, KLG: 4.86, , PSNR: 17.59, SSIM: 0.49,]
Train :: Epoch: 2/10:   3%|▎         | 10/391 [00:08<05:37,  1.13it/s, s=MSE: 0.02, L1: 0.09, , RB: 1725.91, RT: 6.24, , MV: 0.01, KLG: 4.86, , PSNR: 17.59, SSIM: 0.49,]
Train :: Epoch: 2/10:   3%|▎         | 10/391 [00:09<05:37,  1.13it/s, s=MSE: 0.01, L1: 0.09, , RB: 1725.91, RT: 6.24, , MV: 0.01, KLG: 4.85, , PSNR: 17.78, SSIM: 0.48,]
Train :: Epoch: 2/10:   3%|▎         | 11/391 [00:09<05:40,  1.12it/s, s=MSE: 0.01, L1: 0.09, , RB: 1725.91, RT: 6.24, , MV: 0.01, KLG: 4.85, , PSNR: 17.78, SSIM: 0.48,]
Train :: Epoch: 2/10:   3%|▎         | 11/391 [00:09<05:40,  1.12it/s, s=MSE: 0.01, L1: 0.09, , RB: 1725.91, RT: 6.24, , MV: 0.01, KLG: 4.69, , PSNR: 18.21, SSIM: 0.51,]
Train :: Epoch: 2/10:   3%|▎         | 11/391 [00:09<05:44,  1.10it/s, s=MSE: 0.01, L1: 0.09, , RB: 1725.91, RT: 6.24, , MV: 0.01, KLG: 4.69, , PSNR: 18.21, SSIM: 0.51,]

  0%|          | 0/391 [00:00<?, ?it/s]
Train :: Epoch: 3/10:   0%|          | 0/391 [00:00<?, ?it/s]
Train :: Epoch: 3/10:   0%|          | 1/391 [00:00<05:52,  1.11it/s, s=MSE: 0.01, L1: 0.09, , RB: 1725.91, RT: 6.24, , MV: 0.01, KLG: 4.63, , PSNR: 18.12, SSIM: 0.51,]
Train :: Epoch: 3/10:   1%|          | 2/391 [00:00<02:56,  2.21it/s, s=MSE: 0.01, L1: 0.09, , RB: 1725.91, RT: 6.24, , MV: 0.01, KLG: 4.63, , PSNR: 18.12, SSIM: 0.51,]
Train :: Epoch: 3/10:   1%|          | 2/391 [00:01<02:56,  2.21it/s, s=MSE: 0.01, L1: 0.08, , RB: 1725.91, RT: 6.24, , MV: 0.01, KLG: 4.67, , PSNR: 18.46, SSIM: 0.52,]
Train :: Epoch: 3/10:   1%|          | 3/391 [00:01<04:07,  1.57it/s, s=MSE: 0.01, L1: 0.08, , RB: 1725.91, RT: 6.24, , MV: 0.01, KLG: 4.67, , PSNR: 18.46, SSIM: 0.52,]
Train :: Epoch: 3/10:   1%|          | 3/391 [00:02<04:07,  1.57it/s, s=MSE: 0.01, L1: 0.09, , RB: 1725.91, RT: 6.24, , MV: 0.01, KLG: 4.81, , PSNR: 18.20, SSIM: 0.52,]
Train :: Epoch: 3/10:   1%|          | 4/391 [00:02<04:47,  1.35it/s, s=MSE: 0.01, L1: 0.09, , RB: 1725.91, RT: 6.24, , MV: 0.01, KLG: 4.81, , PSNR: 18.20, SSIM: 0.52,]
Train :: Epoch: 3/10:   1%|          | 4/391 [00:03<04:47,  1.35it/s, s=MSE: 0.01, L1: 0.09, , RB: 1725.91, RT: 6.24, , MV: 0.01, KLG: 4.72, , PSNR: 18.06, SSIM: 0.53,]
Train :: Epoch: 3/10:   1%|▏         | 5/391 [00:03<05:07,  1.25it/s, s=MSE: 0.01, L1: 0.09, , RB: 1725.91, RT: 6.24, , MV: 0.01, KLG: 4.72, , PSNR: 18.06, SSIM: 0.53,]
Train :: Epoch: 3/10:   1%|▏         | 5/391 [00:04<05:07,  1.25it/s, s=MSE: 0.01, L1: 0.08, , RB: 1725.91, RT: 6.24, , MV: 0.01, KLG: 4.71, , PSNR: 18.65, SSIM: 0.54,]
Train :: Epoch: 3/10:   2%|▏         | 6/391 [00:04<05:20,  1.20it/s, s=MSE: 0.01, L1: 0.08, , RB: 1725.91, RT: 6.24, , MV: 0.01, KLG: 4.71, , PSNR: 18.65, SSIM: 0.54,]
Train :: Epoch: 3/10:   2%|▏         | 6/391 [00:05<05:20,  1.20it/s, s=MSE: 0.01, L1: 0.08, , RB: 1725.91, RT: 6.24, , MV: 0.01, KLG: 4.88, , PSNR: 18.65, SSIM: 0.55,]
Train :: Epoch: 3/10:   2%|▏         | 7/391 [00:05<05:28,  1.17it/s, s=MSE: 0.01, L1: 0.08, , RB: 1725.91, RT: 6.24, , MV: 0.01, KLG: 4.88, , PSNR: 18.65, SSIM: 0.55,]
Train :: Epoch: 3/10:   2%|▏         | 7/391 [00:06<05:28,  1.17it/s, s=MSE: 0.01, L1: 0.09, , RB: 1725.91, RT: 6.24, , MV: 0.01, KLG: 4.83, , PSNR: 18.27, SSIM: 0.55,]
Train :: Epoch: 3/10:   2%|▏         | 8/391 [00:06<05:32,  1.15it/s, s=MSE: 0.01, L1: 0.09, , RB: 1725.91, RT: 6.24, , MV: 0.01, KLG: 4.83, , PSNR: 18.27, SSIM: 0.55,]
Train :: Epoch: 3/10:   2%|▏         | 8/391 [00:07<05:32,  1.15it/s, s=MSE: 0.01, L1: 0.08, , RB: 1725.91, RT: 6.24, , MV: 0.01, KLG: 4.67, , PSNR: 18.84, SSIM: 0.55,]
Train :: Epoch: 3/10:   2%|▏         | 9/391 [00:07<05:36,  1.14it/s, s=MSE: 0.01, L1: 0.08, , RB: 1725.91, RT: 6.24, , MV: 0.01, KLG: 4.67, , PSNR: 18.84, SSIM: 0.55,]
Train :: Epoch: 3/10:   2%|▏         | 9/391 [00:08<05:36,  1.14it/s, s=MSE: 0.01, L1: 0.08, , RB: 1725.91, RT: 6.24, , MV: 0.01, KLG: 4.86, , PSNR: 18.91, SSIM: 0.57,]
Train :: Epoch: 3/10:   3%|▎         | 10/391 [00:08<05:38,  1.13it/s, s=MSE: 0.01, L1: 0.08, , RB: 1725.91, RT: 6.24, , MV: 0.01, KLG: 4.86, , PSNR: 18.91, SSIM: 0.57,]
Train :: Epoch: 3/10:   3%|▎         | 10/391 [00:09<05:38,  1.13it/s, s=MSE: 0.01, L1: 0.08, , RB: 1725.91, RT: 6.24, , MV: 0.01, KLG: 4.85, , PSNR: 19.03, SSIM: 0.56,]
Train :: Epoch: 3/10:   3%|▎         | 11/391 [00:09<05:39,  1.12it/s, s=MSE: 0.01, L1: 0.08, , RB: 1725.91, RT: 6.24, , MV: 0.01, KLG: 4.85, , PSNR: 19.03, SSIM: 0.56,]
Train :: Epoch: 3/10:   3%|▎         | 11/391 [00:09<05:39,  1.12it/s, s=MSE: 0.01, L1: 0.07, , RB: 1725.91, RT: 6.24, , MV: 0.01, KLG: 4.69, , PSNR: 19.59, SSIM: 0.58,]
Train :: Epoch: 3/10:   3%|▎         | 11/391 [00:09<05:44,  1.10it/s, s=MSE: 0.01, L1: 0.07, , RB: 1725.91, RT: 6.24, , MV: 0.01, KLG: 4.69, , PSNR: 19.59, SSIM: 0.58,]

  0%|          | 0/391 [00:00<?, ?it/s]
Train :: Epoch: 4/10:   0%|          | 0/391 [00:00<?, ?it/s]
Train :: Epoch: 4/10:   0%|          | 1/391 [00:00<05:55,  1.10it/s, s=MSE: 0.01, L1: 0.08, , RB: 1725.91, RT: 6.24, , MV: 0.01, KLG: 4.63, , PSNR: 19.18, SSIM: 0.58,]
Train :: Epoch: 4/10:   1%|          | 2/391 [00:00<02:57,  2.19it/s, s=MSE: 0.01, L1: 0.08, , RB: 1725.91, RT: 6.24, , MV: 0.01, KLG: 4.63, , PSNR: 19.18, SSIM: 0.58,]
Train :: Epoch: 4/10:   1%|          | 2/391 [00:01<02:57,  2.19it/s, s=MSE: 0.01, L1: 0.07, , RB: 1725.91, RT: 6.24, , MV: 0.01, KLG: 4.67, , PSNR: 19.64, SSIM: 0.59,]
Train :: Epoch: 4/10:   1%|          | 3/391 [00:01<04:09,  1.56it/s, s=MSE: 0.01, L1: 0.07, , RB: 1725.91, RT: 6.24, , MV: 0.01, KLG: 4.67, , PSNR: 19.64, SSIM: 0.59,]
Train :: Epoch: 4/10:   1%|          | 3/391 [00:02<04:09,  1.56it/s, s=MSE: 0.01, L1: 0.08, , RB: 1725.91, RT: 6.24, , MV: 0.01, KLG: 4.81, , PSNR: 19.39, SSIM: 0.59,]
Train :: Epoch: 4/10:   1%|          | 4/391 [00:02<04:46,  1.35it/s, s=MSE: 0.01, L1: 0.08, , RB: 1725.91, RT: 6.24, , MV: 0.01, KLG: 4.81, , PSNR: 19.39, SSIM: 0.59,]
Train :: Epoch: 4/10:   1%|          | 4/391 [00:03<04:46,  1.35it/s, s=MSE: 0.01, L1: 0.08, , RB: 1725.91, RT: 6.24, , MV: 0.01, KLG: 4.72, , PSNR: 19.25, SSIM: 0.59,]
Train :: Epoch: 4/10:   1%|▏         | 5/391 [00:03<05:08,  1.25it/s, s=MSE: 0.01, L1: 0.08, , RB: 1725.91, RT: 6.24, , MV: 0.01, KLG: 4.72, , PSNR: 19.25, SSIM: 0.59,]
Train :: Epoch: 4/10:   1%|▏         | 5/391 [00:04<05:08,  1.25it/s, s=MSE: 0.01, L1: 0.07, , RB: 1725.91, RT: 6.24, , MV: 0.01, KLG: 4.71, , PSNR: 19.77, SSIM: 0.60,]
Train :: Epoch: 4/10:   2%|▏         | 6/391 [00:04<05:20,  1.20it/s, s=MSE: 0.01, L1: 0.07, , RB: 1725.91, RT: 6.24, , MV: 0.01, KLG: 4.71, , PSNR: 19.77, SSIM: 0.60,]
Train :: Epoch: 4/10:   2%|▏         | 6/391 [00:05<05:20,  1.20it/s, s=MSE: 0.01, L1: 0.07, , RB: 1725.91, RT: 6.24, , MV: 0.01, KLG: 4.88, , PSNR: 19.78, SSIM: 0.60,]
Train :: Epoch: 4/10:   2%|▏         | 7/391 [00:05<05:29,  1.17it/s, s=MSE: 0.01, L1: 0.07, , RB: 1725.91, RT: 6.24, , MV: 0.01, KLG: 4.88, , PSNR: 19.78, SSIM: 0.60,]
Train :: Epoch: 4/10:   2%|▏         | 7/391 [00:06<05:29,  1.17it/s, s=MSE: 0.01, L1: 0.08, , RB: 1725.91, RT: 6.24, , MV: 0.01, KLG: 4.83, , PSNR: 19.35, SSIM: 0.60,]
Train :: Epoch: 4/10:   2%|▏         | 8/391 [00:06<05:34,  1.15it/s, s=MSE: 0.01, L1: 0.08, , RB: 1725.91, RT: 6.24, , MV: 0.01, KLG: 4.83, , PSNR: 19.35, SSIM: 0.60,]
Train :: Epoch: 4/10:   2%|▏         | 8/391 [00:07<05:34,  1.15it/s, s=MSE: 0.01, L1: 0.07, , RB: 1725.91, RT: 6.24, , MV: 0.01, KLG: 4.67, , PSNR: 19.80, SSIM: 0.61,]
Train :: Epoch: 4/10:   2%|▏         | 9/391 [00:07<05:37,  1.13it/s, s=MSE: 0.01, L1: 0.07, , RB: 1725.91, RT: 6.24, , MV: 0.01, KLG: 4.67, , PSNR: 19.80, SSIM: 0.61,]
Train :: Epoch: 4/10:   2%|▏         | 9/391 [00:08<05:37,  1.13it/s, s=MSE: 0.01, L1: 0.07, , RB: 1725.91, RT: 6.24, , MV: 0.01, KLG: 4.86, , PSNR: 19.94, SSIM: 0.62,]
Train :: Epoch: 4/10:   3%|▎         | 10/391 [00:08<05:39,  1.12it/s, s=MSE: 0.01, L1: 0.07, , RB: 1725.91, RT: 6.24, , MV: 0.01, KLG: 4.86, , PSNR: 19.94, SSIM: 0.62,]
Train :: Epoch: 4/10:   3%|▎         | 10/391 [00:09<05:39,  1.12it/s, s=MSE: 0.01, L1: 0.07, , RB: 1725.91, RT: 6.24, , MV: 0.01, KLG: 4.85, , PSNR: 20.02, SSIM: 0.61,]
Train :: Epoch: 4/10:   3%|▎         | 11/391 [00:09<05:39,  1.12it/s, s=MSE: 0.01, L1: 0.07, , RB: 1725.91, RT: 6.24, , MV: 0.01, KLG: 4.85, , PSNR: 20.02, SSIM: 0.61,]
Train :: Epoch: 4/10:   3%|▎         | 11/391 [00:09<05:39,  1.12it/s, s=MSE: 0.01, L1: 0.07, , RB: 1725.91, RT: 6.24, , MV: 0.01, KLG: 4.69, , PSNR: 20.53, SSIM: 0.63,]
Train :: Epoch: 4/10:   3%|▎         | 11/391 [00:09<05:44,  1.10it/s, s=MSE: 0.01, L1: 0.07, , RB: 1725.91, RT: 6.24, , MV: 0.01, KLG: 4.69, , PSNR: 20.53, SSIM: 0.63,]

  0%|          | 0/391 [00:00<?, ?it/s]
Train :: Epoch: 5/10:   0%|          | 0/391 [00:00<?, ?it/s]
Train :: Epoch: 5/10:   0%|          | 1/391 [00:00<05:50,  1.11it/s, s=MSE: 0.01, L1: 0.07, , RB: 1725.91, RT: 6.24, , MV: 0.01, KLG: 4.63, , PSNR: 19.94, SSIM: 0.62,]
Train :: Epoch: 5/10:   1%|          | 2/391 [00:00<02:55,  2.22it/s, s=MSE: 0.01, L1: 0.07, , RB: 1725.91, RT: 6.24, , MV: 0.01, KLG: 4.63, , PSNR: 19.94, SSIM: 0.62,]
Train :: Epoch: 5/10:   1%|          | 2/391 [00:01<02:55,  2.22it/s, s=MSE: 0.01, L1: 0.07, , RB: 1725.91, RT: 6.24, , MV: 0.01, KLG: 4.67, , PSNR: 20.51, SSIM: 0.63,]
Train :: Epoch: 5/10:   1%|          | 3/391 [00:01<04:07,  1.57it/s, s=MSE: 0.01, L1: 0.07, , RB: 1725.91, RT: 6.24, , MV: 0.01, KLG: 4.67, , PSNR: 20.51, SSIM: 0.63,]
Train :: Epoch: 5/10:   1%|          | 3/391 [00:02<04:07,  1.57it/s, s=MSE: 0.01, L1: 0.07, , RB: 1725.91, RT: 6.24, , MV: 0.01, KLG: 4.81, , PSNR: 20.22, SSIM: 0.63,]
Train :: Epoch: 5/10:   1%|          | 4/391 [00:02<04:46,  1.35it/s, s=MSE: 0.01, L1: 0.07, , RB: 1725.91, RT: 6.24, , MV: 0.01, KLG: 4.81, , PSNR: 20.22, SSIM: 0.63,]
Train :: Epoch: 5/10:   1%|          | 4/391 [00:03<04:46,  1.35it/s, s=MSE: 0.01, L1: 0.07, , RB: 1725.91, RT: 6.24, , MV: 0.01, KLG: 4.72, , PSNR: 20.08, SSIM: 0.64,]
Train :: Epoch: 5/10:   1%|▏         | 5/391 [00:03<05:07,  1.26it/s, s=MSE: 0.01, L1: 0.07, , RB: 1725.91, RT: 6.24, , MV: 0.01, KLG: 4.72, , PSNR: 20.08, SSIM: 0.64,]
Train :: Epoch: 5/10:   1%|▏         | 5/391 [00:04<05:07,  1.26it/s, s=MSE: 0.01, L1: 0.07, , RB: 1725.91, RT: 6.24, , MV: 0.01, KLG: 4.71, , PSNR: 20.61, SSIM: 0.64,]
Train :: Epoch: 5/10:   2%|▏         | 6/391 [00:04<05:19,  1.20it/s, s=MSE: 0.01, L1: 0.07, , RB: 1725.91, RT: 6.24, , MV: 0.01, KLG: 4.71, , PSNR: 20.61, SSIM: 0.64,]
Train :: Epoch: 5/10:   2%|▏         | 6/391 [00:05<05:19,  1.20it/s, s=MSE: 0.01, L1: 0.07, , RB: 1725.91, RT: 6.24, , MV: 0.01, KLG: 4.88, , PSNR: 20.56, SSIM: 0.64,]
Train :: Epoch: 5/10:   2%|▏         | 7/391 [00:05<05:27,  1.17it/s, s=MSE: 0.01, L1: 0.07, , RB: 1725.91, RT: 6.24, , MV: 0.01, KLG: 4.88, , PSNR: 20.56, SSIM: 0.64,]
Train :: Epoch: 5/10:   2%|▏         | 7/391 [00:06<05:27,  1.17it/s, s=MSE: 0.01, L1: 0.07, , RB: 1725.91, RT: 6.24, , MV: 0.01, KLG: 4.83, , PSNR: 20.16, SSIM: 0.64,]
Train :: Epoch: 5/10:   2%|▏         | 8/391 [00:06<05:32,  1.15it/s, s=MSE: 0.01, L1: 0.07, , RB: 1725.91, RT: 6.24, , MV: 0.01, KLG: 4.83, , PSNR: 20.16, SSIM: 0.64,]
Train :: Epoch: 5/10:   2%|▏         | 8/391 [00:07<05:32,  1.15it/s, s=MSE: 0.01, L1: 0.07, , RB: 1725.91, RT: 6.24, , MV: 0.01, KLG: 4.67, , PSNR: 20.57, SSIM: 0.65,]
Train :: Epoch: 5/10:   2%|▏         | 9/391 [00:07<05:35,  1.14it/s, s=MSE: 0.01, L1: 0.07, , RB: 1725.91, RT: 6.24, , MV: 0.01, KLG: 4.67, , PSNR: 20.57, SSIM: 0.65,]
Train :: Epoch: 5/10:   2%|▏         | 9/391 [00:08<05:35,  1.14it/s, s=MSE: 0.01, L1: 0.07, , RB: 1725.91, RT: 6.24, , MV: 0.01, KLG: 4.86, , PSNR: 20.69, SSIM: 0.65,]
Train :: Epoch: 5/10:   3%|▎         | 10/391 [00:08<05:38,  1.13it/s, s=MSE: 0.01, L1: 0.07, , RB: 1725.91, RT: 6.24, , MV: 0.01, KLG: 4.86, , PSNR: 20.69, SSIM: 0.65,]
Train :: Epoch: 5/10:   3%|▎         | 10/391 [00:09<05:38,  1.13it/s, s=MSE: 0.01, L1: 0.07, , RB: 1725.91, RT: 6.24, , MV: 0.01, KLG: 4.85, , PSNR: 20.73, SSIM: 0.64,]
Train :: Epoch: 5/10:   3%|▎         | 11/391 [00:09<05:39,  1.12it/s, s=MSE: 0.01, L1: 0.07, , RB: 1725.91, RT: 6.24, , MV: 0.01, KLG: 4.85, , PSNR: 20.73, SSIM: 0.64,]
Train :: Epoch: 5/10:   3%|▎         | 11/391 [00:09<05:39,  1.12it/s, s=MSE: 0.01, L1: 0.06, , RB: 1725.91, RT: 6.24, , MV: 0.01, KLG: 4.69, , PSNR: 21.20, SSIM: 0.66,]
Train :: Epoch: 5/10:   3%|▎         | 11/391 [00:09<05:43,  1.11it/s, s=MSE: 0.01, L1: 0.06, , RB: 1725.91, RT: 6.24, , MV: 0.01, KLG: 4.69, , PSNR: 21.20, SSIM: 0.66,]

  0%|          | 0/391 [00:00<?, ?it/s]
Train :: Epoch: 6/10:   0%|          | 0/391 [00:00<?, ?it/s]
Train :: Epoch: 6/10:   0%|          | 1/391 [00:00<05:54,  1.10it/s, s=MSE: 0.01, L1: 0.07, , RB: 1725.91, RT: 6.24, , MV: 0.01, KLG: 4.63, , PSNR: 20.53, SSIM: 0.65,]
Train :: Epoch: 6/10:   1%|          | 2/391 [00:00<02:56,  2.20it/s, s=MSE: 0.01, L1: 0.07, , RB: 1725.91, RT: 6.24, , MV: 0.01, KLG: 4.63, , PSNR: 20.53, SSIM: 0.65,]
Train :: Epoch: 6/10:   1%|          | 2/391 [00:01<02:56,  2.20it/s, s=MSE: 0.01, L1: 0.06, , RB: 1725.91, RT: 6.24, , MV: 0.01, KLG: 4.67, , PSNR: 21.02, SSIM: 0.66,]
Train :: Epoch: 6/10:   1%|          | 3/391 [00:01<04:09,  1.55it/s, s=MSE: 0.01, L1: 0.06, , RB: 1725.91, RT: 6.24, , MV: 0.01, KLG: 4.67, , PSNR: 21.02, SSIM: 0.66,]
Train :: Epoch: 6/10:   1%|          | 3/391 [00:02<04:09,  1.55it/s, s=MSE: 0.01, L1: 0.07, , RB: 1725.91, RT: 6.24, , MV: 0.01, KLG: 4.81, , PSNR: 20.71, SSIM: 0.66,]
Train :: Epoch: 6/10:   1%|          | 4/391 [00:02<04:47,  1.35it/s, s=MSE: 0.01, L1: 0.07, , RB: 1725.91, RT: 6.24, , MV: 0.01, KLG: 4.81, , PSNR: 20.71, SSIM: 0.66,]
Train :: Epoch: 6/10:   1%|          | 4/391 [00:03<04:47,  1.35it/s, s=MSE: 0.01, L1: 0.07, , RB: 1725.91, RT: 6.24, , MV: 0.01, KLG: 4.72, , PSNR: 20.64, SSIM: 0.66,]
Train :: Epoch: 6/10:   1%|▏         | 5/391 [00:03<05:08,  1.25it/s, s=MSE: 0.01, L1: 0.07, , RB: 1725.91, RT: 6.24, , MV: 0.01, KLG: 4.72, , PSNR: 20.64, SSIM: 0.66,]
Train :: Epoch: 6/10:   1%|▏         | 5/391 [00:04<05:08,  1.25it/s, s=MSE: 0.01, L1: 0.06, , RB: 1725.91, RT: 6.24, , MV: 0.01, KLG: 4.71, , PSNR: 21.15, SSIM: 0.67,]
Train :: Epoch: 6/10:   2%|▏         | 6/391 [00:04<05:21,  1.20it/s, s=MSE: 0.01, L1: 0.06, , RB: 1725.91, RT: 6.24, , MV: 0.01, KLG: 4.71, , PSNR: 21.15, SSIM: 0.67,]
Train :: Epoch: 6/10:   2%|▏         | 6/391 [00:05<05:21,  1.20it/s, s=MSE: 0.01, L1: 0.06, , RB: 1725.91, RT: 6.24, , MV: 0.01, KLG: 4.88, , PSNR: 21.04, SSIM: 0.67,]
Train :: Epoch: 6/10:   2%|▏         | 7/391 [00:05<05:29,  1.17it/s, s=MSE: 0.01, L1: 0.06, , RB: 1725.91, RT: 6.24, , MV: 0.01, KLG: 4.88, , PSNR: 21.04, SSIM: 0.67,]
Train :: Epoch: 6/10:   2%|▏         | 7/391 [00:06<05:29,  1.17it/s, s=MSE: 0.01, L1: 0.07, , RB: 1725.91, RT: 6.24, , MV: 0.01, KLG: 4.83, , PSNR: 20.75, SSIM: 0.67,]
Train :: Epoch: 6/10:   2%|▏         | 8/391 [00:06<05:34,  1.15it/s, s=MSE: 0.01, L1: 0.07, , RB: 1725.91, RT: 6.24, , MV: 0.01, KLG: 4.83, , PSNR: 20.75, SSIM: 0.67,]
Train :: Epoch: 6/10:   2%|▏         | 8/391 [00:07<05:34,  1.15it/s, s=MSE: 0.01, L1: 0.06, , RB: 1725.91, RT: 6.24, , MV: 0.01, KLG: 4.67, , PSNR: 21.02, SSIM: 0.67,]
Train :: Epoch: 6/10:   2%|▏         | 9/391 [00:07<05:37,  1.13it/s, s=MSE: 0.01, L1: 0.06, , RB: 1725.91, RT: 6.24, , MV: 0.01, KLG: 4.67, , PSNR: 21.02, SSIM: 0.67,]
Train :: Epoch: 6/10:   2%|▏         | 9/391 [00:08<05:37,  1.13it/s, s=MSE: 0.01, L1: 0.06, , RB: 1725.91, RT: 6.24, , MV: 0.01, KLG: 4.86, , PSNR: 21.18, SSIM: 0.68,]
Train :: Epoch: 6/10:   3%|▎         | 10/391 [00:08<05:39,  1.12it/s, s=MSE: 0.01, L1: 0.06, , RB: 1725.91, RT: 6.24, , MV: 0.01, KLG: 4.86, , PSNR: 21.18, SSIM: 0.68,]
Train :: Epoch: 6/10:   3%|▎         | 10/391 [00:09<05:39,  1.12it/s, s=MSE: 0.01, L1: 0.06, , RB: 1725.91, RT: 6.24, , MV: 0.01, KLG: 4.85, , PSNR: 21.16, SSIM: 0.67,]
Train :: Epoch: 6/10:   3%|▎         | 11/391 [00:09<05:39,  1.12it/s, s=MSE: 0.01, L1: 0.06, , RB: 1725.91, RT: 6.24, , MV: 0.01, KLG: 4.85, , PSNR: 21.16, SSIM: 0.67,]
Train :: Epoch: 6/10:   3%|▎         | 11/391 [00:09<05:39,  1.12it/s, s=MSE: 0.01, L1: 0.06, , RB: 1725.91, RT: 6.24, , MV: 0.01, KLG: 4.69, , PSNR: 21.60, SSIM: 0.68,]
Train :: Epoch: 6/10:   3%|▎         | 11/391 [00:09<05:44,  1.10it/s, s=MSE: 0.01, L1: 0.06, , RB: 1725.91, RT: 6.24, , MV: 0.01, KLG: 4.69, , PSNR: 21.60, SSIM: 0.68,]

  0%|          | 0/391 [00:00<?, ?it/s]
Train :: Epoch: 7/10:   0%|          | 0/391 [00:00<?, ?it/s]
Train :: Epoch: 7/10:   0%|          | 1/391 [00:00<05:52,  1.11it/s, s=MSE: 0.01, L1: 0.06, , RB: 1725.91, RT: 6.24, , MV: 0.01, KLG: 4.63, , PSNR: 20.93, SSIM: 0.68,]
Train :: Epoch: 7/10:   1%|          | 2/391 [00:00<02:56,  2.21it/s, s=MSE: 0.01, L1: 0.06, , RB: 1725.91, RT: 6.24, , MV: 0.01, KLG: 4.63, , PSNR: 20.93, SSIM: 0.68,]
Train :: Epoch: 7/10:   1%|          | 2/391 [00:01<02:56,  2.21it/s, s=MSE: 0.01, L1: 0.06, , RB: 1725.91, RT: 6.24, , MV: 0.01, KLG: 4.67, , PSNR: 21.53, SSIM: 0.69,]
Train :: Epoch: 7/10:   1%|          | 3/391 [00:01<04:07,  1.57it/s, s=MSE: 0.01, L1: 0.06, , RB: 1725.91, RT: 6.24, , MV: 0.01, KLG: 4.67, , PSNR: 21.53, SSIM: 0.69,]
Train :: Epoch: 7/10:   1%|          | 3/391 [00:02<04:07,  1.57it/s, s=MSE: 0.01, L1: 0.06, , RB: 1725.91, RT: 6.24, , MV: 0.01, KLG: 4.81, , PSNR: 21.16, SSIM: 0.68,]
Train :: Epoch: 7/10:   1%|          | 4/391 [00:02<04:45,  1.36it/s, s=MSE: 0.01, L1: 0.06, , RB: 1725.91, RT: 6.24, , MV: 0.01, KLG: 4.81, , PSNR: 21.16, SSIM: 0.68,]
Train :: Epoch: 7/10:   1%|          | 4/391 [00:03<04:45,  1.36it/s, s=MSE: 0.01, L1: 0.06, , RB: 1725.91, RT: 6.24, , MV: 0.01, KLG: 4.72, , PSNR: 21.11, SSIM: 0.69,]
Train :: Epoch: 7/10:   1%|▏         | 5/391 [00:03<05:06,  1.26it/s, s=MSE: 0.01, L1: 0.06, , RB: 1725.91, RT: 6.24, , MV: 0.01, KLG: 4.72, , PSNR: 21.11, SSIM: 0.69,]
Train :: Epoch: 7/10:   1%|▏         | 5/391 [00:04<05:06,  1.26it/s, s=MSE: 0.01, L1: 0.06, , RB: 1725.91, RT: 6.24, , MV: 0.01, KLG: 4.71, , PSNR: 21.49, SSIM: 0.70,]
Train :: Epoch: 7/10:   2%|▏         | 6/391 [00:04<05:19,  1.21it/s, s=MSE: 0.01, L1: 0.06, , RB: 1725.91, RT: 6.24, , MV: 0.01, KLG: 4.71, , PSNR: 21.49, SSIM: 0.70,]
Train :: Epoch: 7/10:   2%|▏         | 6/391 [00:05<05:19,  1.21it/s, s=MSE: 0.01, L1: 0.06, , RB: 1725.91, RT: 6.24, , MV: 0.01, KLG: 4.88, , PSNR: 21.45, SSIM: 0.69,]
Train :: Epoch: 7/10:   2%|▏         | 7/391 [00:05<05:27,  1.17it/s, s=MSE: 0.01, L1: 0.06, , RB: 1725.91, RT: 6.24, , MV: 0.01, KLG: 4.88, , PSNR: 21.45, SSIM: 0.69,]
Train :: Epoch: 7/10:   2%|▏         | 7/391 [00:06<05:27,  1.17it/s, s=MSE: 0.01, L1: 0.06, , RB: 1725.91, RT: 6.24, , MV: 0.01, KLG: 4.83, , PSNR: 21.20, SSIM: 0.69,]
Train :: Epoch: 7/10:   2%|▏         | 8/391 [00:06<05:32,  1.15it/s, s=MSE: 0.01, L1: 0.06, , RB: 1725.91, RT: 6.24, , MV: 0.01, KLG: 4.83, , PSNR: 21.20, SSIM: 0.69,]
Train :: Epoch: 7/10:   2%|▏         | 8/391 [00:07<05:32,  1.15it/s, s=MSE: 0.01, L1: 0.06, , RB: 1725.91, RT: 6.24, , MV: 0.01, KLG: 4.67, , PSNR: 21.44, SSIM: 0.70,]
Train :: Epoch: 7/10:   2%|▏         | 9/391 [00:07<05:35,  1.14it/s, s=MSE: 0.01, L1: 0.06, , RB: 1725.91, RT: 6.24, , MV: 0.01, KLG: 4.67, , PSNR: 21.44, SSIM: 0.70,]
Train :: Epoch: 7/10:   2%|▏         | 9/391 [00:08<05:35,  1.14it/s, s=MSE: 0.01, L1: 0.06, , RB: 1725.91, RT: 6.24, , MV: 0.01, KLG: 4.86, , PSNR: 21.61, SSIM: 0.70,]
Train :: Epoch: 7/10:   3%|▎         | 10/391 [00:08<05:37,  1.13it/s, s=MSE: 0.01, L1: 0.06, , RB: 1725.91, RT: 6.24, , MV: 0.01, KLG: 4.86, , PSNR: 21.61, SSIM: 0.70,]
Train :: Epoch: 7/10:   3%|▎         | 10/391 [00:09<05:37,  1.13it/s, s=MSE: 0.01, L1: 0.06, , RB: 1725.91, RT: 6.24, , MV: 0.01, KLG: 4.85, , PSNR: 21.63, SSIM: 0.69,]
Train :: Epoch: 7/10:   3%|▎         | 11/391 [00:09<05:39,  1.12it/s, s=MSE: 0.01, L1: 0.06, , RB: 1725.91, RT: 6.24, , MV: 0.01, KLG: 4.85, , PSNR: 21.63, SSIM: 0.69,]
Train :: Epoch: 7/10:   3%|▎         | 11/391 [00:09<05:39,  1.12it/s, s=MSE: 0.01, L1: 0.06, , RB: 1725.91, RT: 6.24, , MV: 0.01, KLG: 4.69, , PSNR: 21.95, SSIM: 0.70,]
Train :: Epoch: 7/10:   3%|▎         | 11/391 [00:09<05:42,  1.11it/s, s=MSE: 0.01, L1: 0.06, , RB: 1725.91, RT: 6.24, , MV: 0.01, KLG: 4.69, , PSNR: 21.95, SSIM: 0.70,]

  0%|          | 0/391 [00:00<?, ?it/s]
Train :: Epoch: 8/10:   0%|          | 0/391 [00:00<?, ?it/s]
Train :: Epoch: 8/10:   0%|          | 1/391 [00:00<06:04,  1.07it/s, s=MSE: 0.01, L1: 0.06, , RB: 1725.91, RT: 6.24, , MV: 0.01, KLG: 4.63, , PSNR: 21.34, SSIM: 0.70,]
Train :: Epoch: 8/10:   1%|          | 2/391 [00:00<03:02,  2.13it/s, s=MSE: 0.01, L1: 0.06, , RB: 1725.91, RT: 6.24, , MV: 0.01, KLG: 4.63, , PSNR: 21.34, SSIM: 0.70,]
Train :: Epoch: 8/10:   1%|          | 2/391 [00:01<03:02,  2.13it/s, s=MSE: 0.01, L1: 0.06, , RB: 1725.91, RT: 6.24, , MV: 0.01, KLG: 4.67, , PSNR: 21.88, SSIM: 0.70,]
Train :: Epoch: 8/10:   1%|          | 3/391 [00:01<04:12,  1.54it/s, s=MSE: 0.01, L1: 0.06, , RB: 1725.91, RT: 6.24, , MV: 0.01, KLG: 4.67, , PSNR: 21.88, SSIM: 0.70,]
Train :: Epoch: 8/10:   1%|          | 3/391 [00:02<04:12,  1.54it/s, s=MSE: 0.01, L1: 0.06, , RB: 1725.91, RT: 6.24, , MV: 0.01, KLG: 4.81, , PSNR: 21.51, SSIM: 0.70,]
Train :: Epoch: 8/10:   1%|          | 4/391 [00:02<04:49,  1.34it/s, s=MSE: 0.01, L1: 0.06, , RB: 1725.91, RT: 6.24, , MV: 0.01, KLG: 4.81, , PSNR: 21.51, SSIM: 0.70,]
Train :: Epoch: 8/10:   1%|          | 4/391 [00:03<04:49,  1.34it/s, s=MSE: 0.01, L1: 0.06, , RB: 1725.91, RT: 6.24, , MV: 0.01, KLG: 4.72, , PSNR: 21.47, SSIM: 0.71,]
Train :: Epoch: 8/10:   1%|▏         | 5/391 [00:03<05:10,  1.24it/s, s=MSE: 0.01, L1: 0.06, , RB: 1725.91, RT: 6.24, , MV: 0.01, KLG: 4.72, , PSNR: 21.47, SSIM: 0.71,]
Train :: Epoch: 8/10:   1%|▏         | 5/391 [00:04<05:10,  1.24it/s, s=MSE: 0.01, L1: 0.06, , RB: 1725.91, RT: 6.24, , MV: 0.01, KLG: 4.71, , PSNR: 21.80, SSIM: 0.72,]
Train :: Epoch: 8/10:   2%|▏         | 6/391 [00:04<05:23,  1.19it/s, s=MSE: 0.01, L1: 0.06, , RB: 1725.91, RT: 6.24, , MV: 0.01, KLG: 4.71, , PSNR: 21.80, SSIM: 0.72,]
Train :: Epoch: 8/10:   2%|▏         | 6/391 [00:05<05:23,  1.19it/s, s=MSE: 0.01, L1: 0.06, , RB: 1725.91, RT: 6.24, , MV: 0.01, KLG: 4.88, , PSNR: 21.81, SSIM: 0.71,]
Train :: Epoch: 8/10:   2%|▏         | 7/391 [00:05<05:32,  1.16it/s, s=MSE: 0.01, L1: 0.06, , RB: 1725.91, RT: 6.24, , MV: 0.01, KLG: 4.88, , PSNR: 21.81, SSIM: 0.71,]
Train :: Epoch: 8/10:   2%|▏         | 7/391 [00:06<05:32,  1.16it/s, s=MSE: 0.01, L1: 0.06, , RB: 1725.91, RT: 6.24, , MV: 0.01, KLG: 4.83, , PSNR: 21.53, SSIM: 0.71,]
Train :: Epoch: 8/10:   2%|▏         | 8/391 [00:06<05:37,  1.13it/s, s=MSE: 0.01, L1: 0.06, , RB: 1725.91, RT: 6.24, , MV: 0.01, KLG: 4.83, , PSNR: 21.53, SSIM: 0.71,]
Train :: Epoch: 8/10:   2%|▏         | 8/391 [00:07<05:37,  1.13it/s, s=MSE: 0.01, L1: 0.06, , RB: 1725.91, RT: 6.24, , MV: 0.01, KLG: 4.67, , PSNR: 21.74, SSIM: 0.71,]
Train :: Epoch: 8/10:   2%|▏         | 9/391 [00:07<05:41,  1.12it/s, s=MSE: 0.01, L1: 0.06, , RB: 1725.91, RT: 6.24, , MV: 0.01, KLG: 4.67, , PSNR: 21.74, SSIM: 0.71,]
Train :: Epoch: 8/10:   2%|▏         | 9/391 [00:08<05:41,  1.12it/s, s=MSE: 0.01, L1: 0.06, , RB: 1725.91, RT: 6.24, , MV: 0.01, KLG: 4.86, , PSNR: 21.97, SSIM: 0.72,]
Train :: Epoch: 8/10:   3%|▎         | 10/391 [00:08<05:43,  1.11it/s, s=MSE: 0.01, L1: 0.06, , RB: 1725.91, RT: 6.24, , MV: 0.01, KLG: 4.86, , PSNR: 21.97, SSIM: 0.72,]
Train :: Epoch: 8/10:   3%|▎         | 10/391 [00:09<05:43,  1.11it/s, s=MSE: 0.01, L1: 0.06, , RB: 1725.91, RT: 6.24, , MV: 0.01, KLG: 4.85, , PSNR: 21.91, SSIM: 0.71,]
Train :: Epoch: 8/10:   3%|▎         | 11/391 [00:09<05:44,  1.10it/s, s=MSE: 0.01, L1: 0.06, , RB: 1725.91, RT: 6.24, , MV: 0.01, KLG: 4.85, , PSNR: 21.91, SSIM: 0.71,]
Train :: Epoch: 8/10:   3%|▎         | 11/391 [00:10<05:44,  1.10it/s, s=MSE: 0.01, L1: 0.05, , RB: 1725.91, RT: 6.24, , MV: 0.01, KLG: 4.69, , PSNR: 22.26, SSIM: 0.72,]
Train :: Epoch: 8/10:   3%|▎         | 11/391 [00:10<05:48,  1.09it/s, s=MSE: 0.01, L1: 0.05, , RB: 1725.91, RT: 6.24, , MV: 0.01, KLG: 4.69, , PSNR: 22.26, SSIM: 0.72,]

  0%|          | 0/391 [00:00<?, ?it/s]
Train :: Epoch: 9/10:   0%|          | 0/391 [00:00<?, ?it/s]
Train :: Epoch: 9/10:   0%|          | 1/391 [00:00<05:56,  1.09it/s, s=MSE: 0.01, L1: 0.06, , RB: 1725.91, RT: 6.24, , MV: 0.01, KLG: 4.63, , PSNR: 21.59, SSIM: 0.71,]
Train :: Epoch: 9/10:   1%|          | 2/391 [00:00<02:58,  2.18it/s, s=MSE: 0.01, L1: 0.06, , RB: 1725.91, RT: 6.24, , MV: 0.01, KLG: 4.63, , PSNR: 21.59, SSIM: 0.71,]
Train :: Epoch: 9/10:   1%|          | 2/391 [00:01<02:58,  2.18it/s, s=MSE: 0.01, L1: 0.05, , RB: 1725.91, RT: 6.24, , MV: 0.01, KLG: 4.67, , PSNR: 22.15, SSIM: 0.72,]
Train :: Epoch: 9/10:   1%|          | 3/391 [00:01<04:09,  1.56it/s, s=MSE: 0.01, L1: 0.05, , RB: 1725.91, RT: 6.24, , MV: 0.01, KLG: 4.67, , PSNR: 22.15, SSIM: 0.72,]
Train :: Epoch: 9/10:   1%|          | 3/391 [00:02<04:09,  1.56it/s, s=MSE: 0.01, L1: 0.06, , RB: 1725.91, RT: 6.24, , MV: 0.01, KLG: 4.81, , PSNR: 21.80, SSIM: 0.72,]
Train :: Epoch: 9/10:   1%|          | 4/391 [00:02<04:46,  1.35it/s, s=MSE: 0.01, L1: 0.06, , RB: 1725.91, RT: 6.24, , MV: 0.01, KLG: 4.81, , PSNR: 21.80, SSIM: 0.72,]
Train :: Epoch: 9/10:   1%|          | 4/391 [00:03<04:46,  1.35it/s, s=MSE: 0.01, L1: 0.06, , RB: 1725.91, RT: 6.24, , MV: 0.01, KLG: 4.72, , PSNR: 21.77, SSIM: 0.72,]
Train :: Epoch: 9/10:   1%|▏         | 5/391 [00:03<05:07,  1.26it/s, s=MSE: 0.01, L1: 0.06, , RB: 1725.91, RT: 6.24, , MV: 0.01, KLG: 4.72, , PSNR: 21.77, SSIM: 0.72,]
Train :: Epoch: 9/10:   1%|▏         | 5/391 [00:04<05:07,  1.26it/s, s=MSE: 0.01, L1: 0.06, , RB: 1725.91, RT: 6.24, , MV: 0.01, KLG: 4.71, , PSNR: 22.08, SSIM: 0.73,]
Train :: Epoch: 9/10:   2%|▏         | 6/391 [00:04<05:20,  1.20it/s, s=MSE: 0.01, L1: 0.06, , RB: 1725.91, RT: 6.24, , MV: 0.01, KLG: 4.71, , PSNR: 22.08, SSIM: 0.73,]
Train :: Epoch: 9/10:   2%|▏         | 6/391 [00:05<05:20,  1.20it/s, s=MSE: 0.01, L1: 0.06, , RB: 1725.91, RT: 6.24, , MV: 0.01, KLG: 4.88, , PSNR: 22.03, SSIM: 0.72,]
Train :: Epoch: 9/10:   2%|▏         | 7/391 [00:05<05:29,  1.17it/s, s=MSE: 0.01, L1: 0.06, , RB: 1725.91, RT: 6.24, , MV: 0.01, KLG: 4.88, , PSNR: 22.03, SSIM: 0.72,]
Train :: Epoch: 9/10:   2%|▏         | 7/391 [00:06<05:29,  1.17it/s, s=MSE: 0.01, L1: 0.06, , RB: 1725.91, RT: 6.24, , MV: 0.01, KLG: 4.83, , PSNR: 21.87, SSIM: 0.72,]
Train :: Epoch: 9/10:   2%|▏         | 8/391 [00:06<05:34,  1.14it/s, s=MSE: 0.01, L1: 0.06, , RB: 1725.91, RT: 6.24, , MV: 0.01, KLG: 4.83, , PSNR: 21.87, SSIM: 0.72,]
Train :: Epoch: 9/10:   2%|▏         | 8/391 [00:07<05:34,  1.14it/s, s=MSE: 0.01, L1: 0.06, , RB: 1725.91, RT: 6.24, , MV: 0.01, KLG: 4.67, , PSNR: 21.98, SSIM: 0.73,]
Train :: Epoch: 9/10:   2%|▏         | 9/391 [00:07<05:37,  1.13it/s, s=MSE: 0.01, L1: 0.06, , RB: 1725.91, RT: 6.24, , MV: 0.01, KLG: 4.67, , PSNR: 21.98, SSIM: 0.73,]
Train :: Epoch: 9/10:   2%|▏         | 9/391 [00:08<05:37,  1.13it/s, s=MSE: 0.01, L1: 0.05, , RB: 1725.91, RT: 6.24, , MV: 0.01, KLG: 4.86, , PSNR: 22.26, SSIM: 0.73,]
Train :: Epoch: 9/10:   3%|▎         | 10/391 [00:08<05:39,  1.12it/s, s=MSE: 0.01, L1: 0.05, , RB: 1725.91, RT: 6.24, , MV: 0.01, KLG: 4.86, , PSNR: 22.26, SSIM: 0.73,]
Train :: Epoch: 9/10:   3%|▎         | 10/391 [00:09<05:39,  1.12it/s, s=MSE: 0.01, L1: 0.06, , RB: 1725.91, RT: 6.24, , MV: 0.01, KLG: 4.85, , PSNR: 22.12, SSIM: 0.72,]
Train :: Epoch: 9/10:   3%|▎         | 11/391 [00:09<05:40,  1.11it/s, s=MSE: 0.01, L1: 0.06, , RB: 1725.91, RT: 6.24, , MV: 0.01, KLG: 4.85, , PSNR: 22.12, SSIM: 0.72,]
Train :: Epoch: 9/10:   3%|▎         | 11/391 [00:09<05:40,  1.11it/s, s=MSE: 0.01, L1: 0.05, , RB: 1725.91, RT: 6.24, , MV: 0.01, KLG: 4.69, , PSNR: 22.54, SSIM: 0.73,]
Train :: Epoch: 9/10:   3%|▎         | 11/391 [00:09<05:44,  1.10it/s, s=MSE: 0.01, L1: 0.05, , RB: 1725.91, RT: 6.24, , MV: 0.01, KLG: 4.69, , PSNR: 22.54, SSIM: 0.73,]

  0%|          | 0/391 [00:00<?, ?it/s]
Train :: Epoch: 10/10:   0%|          | 0/391 [00:00<?, ?it/s]
Train :: Epoch: 10/10:   0%|          | 1/391 [00:00<05:55,  1.10it/s, s=MSE: 0.01, L1: 0.06, , RB: 1725.91, RT: 6.24, , MV: 0.01, KLG: 4.63, , PSNR: 21.79, SSIM: 0.72,]
Train :: Epoch: 10/10:   1%|          | 2/391 [00:00<02:57,  2.19it/s, s=MSE: 0.01, L1: 0.06, , RB: 1725.91, RT: 6.24, , MV: 0.01, KLG: 4.63, , PSNR: 21.79, SSIM: 0.72,]
Train :: Epoch: 10/10:   1%|          | 2/391 [00:01<02:57,  2.19it/s, s=MSE: 0.01, L1: 0.05, , RB: 1725.91, RT: 6.24, , MV: 0.01, KLG: 4.67, , PSNR: 22.41, SSIM: 0.73,]
Train :: Epoch: 10/10:   1%|          | 3/391 [00:01<04:10,  1.55it/s, s=MSE: 0.01, L1: 0.05, , RB: 1725.91, RT: 6.24, , MV: 0.01, KLG: 4.67, , PSNR: 22.41, SSIM: 0.73,]
Train :: Epoch: 10/10:   1%|          | 3/391 [00:02<04:10,  1.55it/s, s=MSE: 0.01, L1: 0.06, , RB: 1725.91, RT: 6.24, , MV: 0.01, KLG: 4.81, , PSNR: 22.00, SSIM: 0.72,]
Train :: Epoch: 10/10:   1%|          | 4/391 [00:02<04:48,  1.34it/s, s=MSE: 0.01, L1: 0.06, , RB: 1725.91, RT: 6.24, , MV: 0.01, KLG: 4.81, , PSNR: 22.00, SSIM: 0.72,]
Train :: Epoch: 10/10:   1%|          | 4/391 [00:03<04:48,  1.34it/s, s=MSE: 0.01, L1: 0.06, , RB: 1725.91, RT: 6.24, , MV: 0.01, KLG: 4.72, , PSNR: 22.04, SSIM: 0.74,]
Train :: Epoch: 10/10:   1%|▏         | 5/391 [00:03<05:09,  1.25it/s, s=MSE: 0.01, L1: 0.06, , RB: 1725.91, RT: 6.24, , MV: 0.01, KLG: 4.72, , PSNR: 22.04, SSIM: 0.74,]
Train :: Epoch: 10/10:   1%|▏         | 5/391 [00:04<05:09,  1.25it/s, s=MSE: 0.01, L1: 0.05, , RB: 1725.91, RT: 6.24, , MV: 0.01, KLG: 4.71, , PSNR: 22.25, SSIM: 0.74,]
Train :: Epoch: 10/10:   2%|▏         | 6/391 [00:04<05:22,  1.19it/s, s=MSE: 0.01, L1: 0.05, , RB: 1725.91, RT: 6.24, , MV: 0.01, KLG: 4.71, , PSNR: 22.25, SSIM: 0.74,]
Train :: Epoch: 10/10:   2%|▏         | 6/391 [00:05<05:22,  1.19it/s, s=MSE: 0.01, L1: 0.05, , RB: 1725.91, RT: 6.24, , MV: 0.01, KLG: 4.88, , PSNR: 22.37, SSIM: 0.74,]
Train :: Epoch: 10/10:   2%|▏         | 7/391 [00:05<05:29,  1.16it/s, s=MSE: 0.01, L1: 0.05, , RB: 1725.91, RT: 6.24, , MV: 0.01, KLG: 4.88, , PSNR: 22.37, SSIM: 0.74,]
Train :: Epoch: 10/10:   2%|▏         | 7/391 [00:06<05:29,  1.16it/s, s=MSE: 0.01, L1: 0.06, , RB: 1725.91, RT: 6.24, , MV: 0.01, KLG: 4.83, , PSNR: 22.01, SSIM: 0.73,]
Train :: Epoch: 10/10:   2%|▏         | 8/391 [00:06<05:34,  1.14it/s, s=MSE: 0.01, L1: 0.06, , RB: 1725.91, RT: 6.24, , MV: 0.01, KLG: 4.83, , PSNR: 22.01, SSIM: 0.73,]
Train :: Epoch: 10/10:   2%|▏         | 8/391 [00:07<05:34,  1.14it/s, s=MSE: 0.01, L1: 0.05, , RB: 1725.91, RT: 6.24, , MV: 0.01, KLG: 4.67, , PSNR: 22.25, SSIM: 0.73,]
Train :: Epoch: 10/10:   2%|▏         | 9/391 [00:07<05:37,  1.13it/s, s=MSE: 0.01, L1: 0.05, , RB: 1725.91, RT: 6.24, , MV: 0.01, KLG: 4.67, , PSNR: 22.25, SSIM: 0.73,]
Train :: Epoch: 10/10:   2%|▏         | 9/391 [00:08<05:37,  1.13it/s, s=MSE: 0.01, L1: 0.05, , RB: 1725.91, RT: 6.24, , MV: 0.01, KLG: 4.86, , PSNR: 22.37, SSIM: 0.74,]
Train :: Epoch: 10/10:   3%|▎         | 10/391 [00:08<05:42,  1.11it/s, s=MSE: 0.01, L1: 0.05, , RB: 1725.91, RT: 6.24, , MV: 0.01, KLG: 4.86, , PSNR: 22.37, SSIM: 0.74,]
Train :: Epoch: 10/10:   3%|▎         | 10/391 [00:09<05:42,  1.11it/s, s=MSE: 0.01, L1: 0.05, , RB: 1725.91, RT: 6.24, , MV: 0.01, KLG: 4.85, , PSNR: 22.39, SSIM: 0.73,]
Train :: Epoch: 10/10:   3%|▎         | 11/391 [00:09<05:41,  1.11it/s, s=MSE: 0.01, L1: 0.05, , RB: 1725.91, RT: 6.24, , MV: 0.01, KLG: 4.85, , PSNR: 22.39, SSIM: 0.73,]
Train :: Epoch: 10/10:   3%|▎         | 11/391 [00:09<05:41,  1.11it/s, s=MSE: 0.01, L1: 0.05, , RB: 1725.91, RT: 6.24, , MV: 0.01, KLG: 4.69, , PSNR: 22.64, SSIM: 0.74,]
Train :: Epoch: 10/10:   3%|▎         | 11/391 [00:09<05:45,  1.10it/s, s=MSE: 0.01, L1: 0.05, , RB: 1725.91, RT: 6.24, , MV: 0.01, KLG: 4.69, , PSNR: 22.64, SSIM: 0.74,]

Visualize results

Performs the inference \(\tilde{\mathbf{x}} = \mathcal{G}_{\theta^*}( \mathbf{H}_{\phi^*}\mathbf{x})\) and visualize the results.

x_est = model(sample.to(device)).cpu()
y = acquisition_model(sample.to(device)).cpu()

normalize = lambda x: (x - torch.min(x)) / (torch.max(x) - torch.min(x))

if acquisition_name == 'spc':
    y = y.reshape(y.shape[0], -1, n_measurements_sqrt, n_measurements_sqrt)

img = make_grid(sample[:16], nrow=4, padding=1, normalize=True, scale_each=False, pad_value=0)
img_est = make_grid(x_est[:16], nrow=4, padding=1, normalize=True, scale_each=False, pad_value=0)
img_y = make_grid(y[:16], nrow=4, padding=1, normalize=True, scale_each=False, pad_value=0)

imgs_dict = {
    "CIFAR10 dataset": img,
    f"{acquisition_name.upper()} measurements": img_y,
    "Recons CIFAR10": img_est
}

plt.figure(figsize=(14, 2.7))

for i, (title, img) in enumerate(imgs_dict.items()):
    plt.subplot(1, 4, i + 1)
    plt.imshow(img.permute(1, 2, 0))
    plt.title(title)
    plt.axis('off')

ca = normalize(acquisition_model.learnable_optics.cpu().detach())
ca = ca.numpy().squeeze()

if acquisition_name == 'spc':
    ca = ca = ca.reshape(n_measurements, 32, 32, 1)[0]
elif acquisition_name == 'c_cassi':
    ca = ca.transpose(1, 2, 0)
plt.subplot(1, 4, 4)
plt.imshow(ca, cmap='gray')
plt.axis('off')
plt.title('Learned CA')
plt.colorbar()

plt.show()
CIFAR10 dataset, C_CASSI measurements, Recons CIFAR10, Learned CA

Total running time of the script: (1 minutes 41.749 seconds)

Gallery generated by Sphinx-Gallery