Note
Go to the end to download the full example code.
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,
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,
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,
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()

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()

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()

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