Simple UNet model.

Contents

Simple UNet model.#

This model is a simplified version of the U-Net architecture, which is widely used for image segmentation tasks. This is implemented in the proprietary FASTMRI package [fastmri].

The U-Net model consists of an encoder (downsampling path) and a decoder (upsampling path) with skip connections between corresponding layers in the encoder and decoder. These skip connections help in retaining spatial information that is lost during the downsampling process.

The primary purpose of this model is to perform image reconstruction tasks, specifically for MRI images. It takes an input MRI image and reconstructs it to improve the image quality or to recover missing parts of the image.

This implementation of the UNet model was pulled from the FastMRI Facebook repository, which is a collaborative research project aimed at advancing the field of medical imaging using machine learning techniques.

\[\mathbf{\hat{x}} = \mathrm{arg} \min_{\mathbf{x}} || \mathcal{U}_\mathbf{\theta}(\mathbf{y}) - \mathbf{x} ||_2^2\]

where \(\mathbf{\hat{x}}\) is the reconstructed MRI image, \(\mathbf{x}\) is the ground truth image, \(\mathbf{y}\) is the input MRI image (e.g., k-space data), and \(\mathcal{U}_\mathbf{\theta}\) is the U-Net model parameterized by \(\theta\).

Warning

We train on a single image here. In practice, this should be done on a database like fastMRI [fastmri].

Imports

import os
from pathlib import Path
import shutil
import brainweb_dl as bwdl
import matplotlib.pyplot as plt
import numpy as np
import torch
from tqdm import tqdm
import time
import joblib
from PIL import Image
import tempfile as tmp

from fastmri.models import Unet
from mrinufft import get_operator
from mrinufft.trajectories import initialize_2D_cones

Setup a simple class for the U-Net model

class Model(torch.nn.Module):
    """Model for MRI reconstruction using a U-Net."""

    def __init__(self, initial_trajectory):
        super().__init__()
        self.operator = get_operator("gpunufft", wrt_data=True)(
            initial_trajectory,
            shape=(256, 256),
            density=True,
            squeeze_dims=False,
        )
        self.unet = Unet(in_chans=1, out_chans=1, chans=32, num_pool_layers=4)

    def forward(self, kspace):
        """Forward pass of the model."""
        image = self.operator.adj_op(kspace)
        recon = self.unet(image.float()).abs()
        recon /= torch.mean(recon)
        return recon

Utility function to plot the state of the model

def plot_state(axs, mri_2D, traj, recon, loss=None, save_name=None):
    """Image plotting function.

    Plot the original MRI image, the trajectory, the reconstructed image,
    and the loss curve (if provided). Saves the plot if a filename is provided.

    Parameters
    ----------
    axs (numpy array): Array of matplotlib axes to plot on.
    mri_2D (torch.Tensor): Original MRI image.
    traj : Trajectory.
    recon (torch.Tensor): Reconstructed image after training.
    loss (list, optional): List of loss values to plot. Defaults to None.
    save_name (str, optional): Filename to save the plot. Defaults to None.
    """
    axs = axs.flatten()
    axs[0].imshow(np.abs(mri_2D[0]), cmap="gray")
    axs[0].axis("off")
    axs[0].set_title("MR Image")
    axs[1].scatter(*traj.T, s=0.5)
    axs[1].set_title("Trajectory")
    axs[2].imshow(np.abs(recon[0][0].detach().cpu().numpy()), cmap="gray")
    axs[2].axis("off")
    axs[2].set_title("Reconstruction")
    if loss is not None:
        axs[3].plot(loss)
        axs[3].grid("on")
        axs[3].set_title("Loss")
    if save_name is not None:
        plt.savefig(save_name, bbox_inches="tight")
        plt.close()
    else:
        plt.show()

Setup Inputs (models, trajectory and image)

init_traj = initialize_2D_cones(32, 256).reshape(-1, 2).astype(np.float32)
model = Model(init_traj)
model.eval()
Model(
  (operator): MRINufftAutoGrad()
  (unet): Unet(
    (down_sample_layers): ModuleList(
      (0): ConvBlock(
        (layers): Sequential(
          (0): Conv2d(1, 32, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)
          (1): InstanceNorm2d(32, eps=1e-05, momentum=0.1, affine=False, track_running_stats=False)
          (2): LeakyReLU(negative_slope=0.2, inplace=True)
          (3): Dropout2d(p=0.0, inplace=False)
          (4): Conv2d(32, 32, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)
          (5): InstanceNorm2d(32, eps=1e-05, momentum=0.1, affine=False, track_running_stats=False)
          (6): LeakyReLU(negative_slope=0.2, inplace=True)
          (7): Dropout2d(p=0.0, inplace=False)
        )
      )
      (1): ConvBlock(
        (layers): Sequential(
          (0): Conv2d(32, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)
          (1): InstanceNorm2d(64, eps=1e-05, momentum=0.1, affine=False, track_running_stats=False)
          (2): LeakyReLU(negative_slope=0.2, inplace=True)
          (3): Dropout2d(p=0.0, inplace=False)
          (4): Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)
          (5): InstanceNorm2d(64, eps=1e-05, momentum=0.1, affine=False, track_running_stats=False)
          (6): LeakyReLU(negative_slope=0.2, inplace=True)
          (7): Dropout2d(p=0.0, inplace=False)
        )
      )
      (2): ConvBlock(
        (layers): Sequential(
          (0): Conv2d(64, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)
          (1): InstanceNorm2d(128, eps=1e-05, momentum=0.1, affine=False, track_running_stats=False)
          (2): LeakyReLU(negative_slope=0.2, inplace=True)
          (3): Dropout2d(p=0.0, inplace=False)
          (4): Conv2d(128, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)
          (5): InstanceNorm2d(128, eps=1e-05, momentum=0.1, affine=False, track_running_stats=False)
          (6): LeakyReLU(negative_slope=0.2, inplace=True)
          (7): Dropout2d(p=0.0, inplace=False)
        )
      )
      (3): ConvBlock(
        (layers): Sequential(
          (0): Conv2d(128, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)
          (1): InstanceNorm2d(256, eps=1e-05, momentum=0.1, affine=False, track_running_stats=False)
          (2): LeakyReLU(negative_slope=0.2, inplace=True)
          (3): Dropout2d(p=0.0, inplace=False)
          (4): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)
          (5): InstanceNorm2d(256, eps=1e-05, momentum=0.1, affine=False, track_running_stats=False)
          (6): LeakyReLU(negative_slope=0.2, inplace=True)
          (7): Dropout2d(p=0.0, inplace=False)
        )
      )
    )
    (conv): ConvBlock(
      (layers): Sequential(
        (0): Conv2d(256, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)
        (1): InstanceNorm2d(512, eps=1e-05, momentum=0.1, affine=False, track_running_stats=False)
        (2): LeakyReLU(negative_slope=0.2, inplace=True)
        (3): Dropout2d(p=0.0, inplace=False)
        (4): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)
        (5): InstanceNorm2d(512, eps=1e-05, momentum=0.1, affine=False, track_running_stats=False)
        (6): LeakyReLU(negative_slope=0.2, inplace=True)
        (7): Dropout2d(p=0.0, inplace=False)
      )
    )
    (up_conv): ModuleList(
      (0): ConvBlock(
        (layers): Sequential(
          (0): Conv2d(512, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)
          (1): InstanceNorm2d(256, eps=1e-05, momentum=0.1, affine=False, track_running_stats=False)
          (2): LeakyReLU(negative_slope=0.2, inplace=True)
          (3): Dropout2d(p=0.0, inplace=False)
          (4): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)
          (5): InstanceNorm2d(256, eps=1e-05, momentum=0.1, affine=False, track_running_stats=False)
          (6): LeakyReLU(negative_slope=0.2, inplace=True)
          (7): Dropout2d(p=0.0, inplace=False)
        )
      )
      (1): ConvBlock(
        (layers): Sequential(
          (0): Conv2d(256, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)
          (1): InstanceNorm2d(128, eps=1e-05, momentum=0.1, affine=False, track_running_stats=False)
          (2): LeakyReLU(negative_slope=0.2, inplace=True)
          (3): Dropout2d(p=0.0, inplace=False)
          (4): Conv2d(128, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)
          (5): InstanceNorm2d(128, eps=1e-05, momentum=0.1, affine=False, track_running_stats=False)
          (6): LeakyReLU(negative_slope=0.2, inplace=True)
          (7): Dropout2d(p=0.0, inplace=False)
        )
      )
      (2): ConvBlock(
        (layers): Sequential(
          (0): Conv2d(128, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)
          (1): InstanceNorm2d(64, eps=1e-05, momentum=0.1, affine=False, track_running_stats=False)
          (2): LeakyReLU(negative_slope=0.2, inplace=True)
          (3): Dropout2d(p=0.0, inplace=False)
          (4): Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)
          (5): InstanceNorm2d(64, eps=1e-05, momentum=0.1, affine=False, track_running_stats=False)
          (6): LeakyReLU(negative_slope=0.2, inplace=True)
          (7): Dropout2d(p=0.0, inplace=False)
        )
      )
      (3): Sequential(
        (0): ConvBlock(
          (layers): Sequential(
            (0): Conv2d(64, 32, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)
            (1): InstanceNorm2d(32, eps=1e-05, momentum=0.1, affine=False, track_running_stats=False)
            (2): LeakyReLU(negative_slope=0.2, inplace=True)
            (3): Dropout2d(p=0.0, inplace=False)
            (4): Conv2d(32, 32, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)
            (5): InstanceNorm2d(32, eps=1e-05, momentum=0.1, affine=False, track_running_stats=False)
            (6): LeakyReLU(negative_slope=0.2, inplace=True)
            (7): Dropout2d(p=0.0, inplace=False)
          )
        )
        (1): Conv2d(32, 1, kernel_size=(1, 1), stride=(1, 1))
      )
    )
    (up_transpose_conv): ModuleList(
      (0): TransposeConvBlock(
        (layers): Sequential(
          (0): ConvTranspose2d(512, 256, kernel_size=(2, 2), stride=(2, 2), bias=False)
          (1): InstanceNorm2d(256, eps=1e-05, momentum=0.1, affine=False, track_running_stats=False)
          (2): LeakyReLU(negative_slope=0.2, inplace=True)
        )
      )
      (1): TransposeConvBlock(
        (layers): Sequential(
          (0): ConvTranspose2d(256, 128, kernel_size=(2, 2), stride=(2, 2), bias=False)
          (1): InstanceNorm2d(128, eps=1e-05, momentum=0.1, affine=False, track_running_stats=False)
          (2): LeakyReLU(negative_slope=0.2, inplace=True)
        )
      )
      (2): TransposeConvBlock(
        (layers): Sequential(
          (0): ConvTranspose2d(128, 64, kernel_size=(2, 2), stride=(2, 2), bias=False)
          (1): InstanceNorm2d(64, eps=1e-05, momentum=0.1, affine=False, track_running_stats=False)
          (2): LeakyReLU(negative_slope=0.2, inplace=True)
        )
      )
      (3): TransposeConvBlock(
        (layers): Sequential(
          (0): ConvTranspose2d(64, 32, kernel_size=(2, 2), stride=(2, 2), bias=False)
          (1): InstanceNorm2d(32, eps=1e-05, momentum=0.1, affine=False, track_running_stats=False)
          (2): LeakyReLU(negative_slope=0.2, inplace=True)
        )
      )
    )
  )
)

Get the image on which we will train our U-Net Model

mri_2D = torch.Tensor(np.flipud(bwdl.get_mri(4, "T1")[80, ...]).astype(np.complex64))[
    None
]
mri_2D = mri_2D / torch.mean(mri_2D)
kspace_mri_2D = model.operator.op(mri_2D)

# Before training, here is the simple reconstruction we have using a
# density compensated adjoint.
dc_adjoint = model.operator.adj_op(kspace_mri_2D)
fig, axs = plt.subplots(1, 3, figsize=(15, 5))
plot_state(axs, mri_2D, init_traj, dc_adjoint)
MR Image, Trajectory, Reconstruction
/volatile/github-ci-mind-inria/gpu_runner/_work/mri-nufft/mri-nufft/examples/GPU/example_fastMRI_UNet.py:131: UserWarning: Casting complex values to real discards the imaginary part (Triggered internally at ../aten/src/ATen/native/Copy.cpp:308.)
  mri_2D = torch.Tensor(np.flipud(bwdl.get_mri(4, "T1")[80, ...]).astype(np.complex64))[

Start training loop

num_epochs = 100
optimizer = torch.optim.RAdam(model.parameters(), lr=1e-3)
losses = []  # Store the loss values and create an animation
image_files = []  # Store the images to create a gif
model.train()

with tqdm(range(num_epochs), unit="steps") as tqdms:
    for i in tqdms:
        out = model(kspace_mri_2D)  # Forward pass

        loss = torch.nn.functional.l1_loss(out, mri_2D[None])  # Compute loss
        tqdms.set_postfix({"loss": loss.item()})  # Update progress bar
        losses.append(loss.item())  # Store loss value

        optimizer.zero_grad()  # Zero gradients
        loss.backward()  # Backward pass
        optimizer.step()  # Update weights

        # Generate images for gif
        hashed = joblib.hash((i, "learn_traj", time.time()))
        filename = f"{tmp.NamedTemporaryFile().name}.png"
        fig, axs = plt.subplots(2, 2, figsize=(10, 10))
        plot_state(
            axs,
            mri_2D,
            init_traj,
            out,
            losses,
            save_name=filename,
        )
        image_files.append(filename)


# Make a GIF of all images.
imgs = [Image.open(img) for img in image_files]
imgs[0].save(
    "mrinufft_learn_unet.gif",
    save_all=True,
    append_images=imgs[1:],
    optimize=False,
    duration=2,
    loop=0,
)

# sphinx_gallery_thumbnail_path = 'generated/autoexamples/GPU/images/mrinufft_learn_unet.gif'
  0%|          | 0/100 [00:00<?, ?steps/s]
  0%|          | 0/100 [00:00<?, ?steps/s, loss=0.706]
  1%|          | 1/100 [00:01<02:30,  1.52s/steps, loss=0.706]
  1%|          | 1/100 [00:01<02:30,  1.52s/steps, loss=0.673]
  2%|▏         | 2/100 [00:02<02:12,  1.35s/steps, loss=0.673]
  2%|▏         | 2/100 [00:03<02:12,  1.35s/steps, loss=0.653]
  3%|▎         | 3/100 [00:04<02:12,  1.37s/steps, loss=0.653]
  3%|▎         | 3/100 [00:04<02:12,  1.37s/steps, loss=0.638]
  4%|▍         | 4/100 [00:05<02:14,  1.40s/steps, loss=0.638]
  4%|▍         | 4/100 [00:05<02:14,  1.40s/steps, loss=0.625]
  5%|▌         | 5/100 [00:06<02:09,  1.36s/steps, loss=0.625]
  5%|▌         | 5/100 [00:07<02:09,  1.36s/steps, loss=0.614]
  6%|▌         | 6/100 [00:08<02:13,  1.42s/steps, loss=0.614]
  6%|▌         | 6/100 [00:08<02:13,  1.42s/steps, loss=0.584]
  7%|▋         | 7/100 [00:09<02:04,  1.34s/steps, loss=0.584]
  7%|▋         | 7/100 [00:09<02:04,  1.34s/steps, loss=0.558]
  8%|▊         | 8/100 [00:10<01:56,  1.27s/steps, loss=0.558]
  8%|▊         | 8/100 [00:10<01:56,  1.27s/steps, loss=0.535]
  9%|▉         | 9/100 [00:11<01:50,  1.21s/steps, loss=0.535]
  9%|▉         | 9/100 [00:11<01:50,  1.21s/steps, loss=0.512]
 10%|█         | 10/100 [00:12<01:45,  1.17s/steps, loss=0.512]
 10%|█         | 10/100 [00:12<01:45,  1.17s/steps, loss=0.49]
 11%|█         | 11/100 [00:13<01:39,  1.12s/steps, loss=0.49]
 11%|█         | 11/100 [00:14<01:39,  1.12s/steps, loss=0.469]
 12%|█▏        | 12/100 [00:14<01:37,  1.11s/steps, loss=0.469]
 12%|█▏        | 12/100 [00:15<01:37,  1.11s/steps, loss=0.449]
 13%|█▎        | 13/100 [00:15<01:30,  1.04s/steps, loss=0.449]
 13%|█▎        | 13/100 [00:15<01:30,  1.04s/steps, loss=0.429]
 14%|█▍        | 14/100 [00:16<01:28,  1.03s/steps, loss=0.429]
 14%|█▍        | 14/100 [00:17<01:28,  1.03s/steps, loss=0.411]
 15%|█▌        | 15/100 [00:17<01:26,  1.02s/steps, loss=0.411]
 15%|█▌        | 15/100 [00:17<01:26,  1.02s/steps, loss=0.393]
 16%|█▌        | 16/100 [00:18<01:28,  1.05s/steps, loss=0.393]
 16%|█▌        | 16/100 [00:19<01:28,  1.05s/steps, loss=0.375]
 17%|█▋        | 17/100 [00:19<01:24,  1.02s/steps, loss=0.375]
 17%|█▋        | 17/100 [00:20<01:24,  1.02s/steps, loss=0.358]
 18%|█▊        | 18/100 [00:20<01:20,  1.02steps/s, loss=0.358]
 18%|█▊        | 18/100 [00:20<01:20,  1.02steps/s, loss=0.341]
 19%|█▉        | 19/100 [00:21<01:21,  1.01s/steps, loss=0.341]
 19%|█▉        | 19/100 [00:22<01:21,  1.01s/steps, loss=0.325]
 20%|██        | 20/100 [00:22<01:17,  1.04steps/s, loss=0.325]
 20%|██        | 20/100 [00:22<01:17,  1.04steps/s, loss=0.31]
 21%|██        | 21/100 [00:23<01:14,  1.07steps/s, loss=0.31]
 21%|██        | 21/100 [00:23<01:14,  1.07steps/s, loss=0.295]
 22%|██▏       | 22/100 [00:24<01:10,  1.11steps/s, loss=0.295]
 22%|██▏       | 22/100 [00:24<01:10,  1.11steps/s, loss=0.28]
 23%|██▎       | 23/100 [00:25<01:09,  1.11steps/s, loss=0.28]
 23%|██▎       | 23/100 [00:25<01:09,  1.11steps/s, loss=0.267]
 24%|██▍       | 24/100 [00:26<01:07,  1.13steps/s, loss=0.267]
 24%|██▍       | 24/100 [00:26<01:07,  1.13steps/s, loss=0.254]
 25%|██▌       | 25/100 [00:27<01:05,  1.14steps/s, loss=0.254]
 25%|██▌       | 25/100 [00:27<01:05,  1.14steps/s, loss=0.242]
 26%|██▌       | 26/100 [00:28<01:09,  1.06steps/s, loss=0.242]
 26%|██▌       | 26/100 [00:28<01:09,  1.06steps/s, loss=0.231]
 27%|██▋       | 27/100 [00:28<01:06,  1.10steps/s, loss=0.231]
 27%|██▋       | 27/100 [00:29<01:06,  1.10steps/s, loss=0.221]
 28%|██▊       | 28/100 [00:29<01:04,  1.12steps/s, loss=0.221]
 28%|██▊       | 28/100 [00:29<01:04,  1.12steps/s, loss=0.211]
 29%|██▉       | 29/100 [00:30<01:02,  1.14steps/s, loss=0.211]
 29%|██▉       | 29/100 [00:30<01:02,  1.14steps/s, loss=0.201]
 30%|███       | 30/100 [00:31<01:01,  1.14steps/s, loss=0.201]
 30%|███       | 30/100 [00:31<01:01,  1.14steps/s, loss=0.192]
 31%|███       | 31/100 [00:32<01:00,  1.14steps/s, loss=0.192]
 31%|███       | 31/100 [00:32<01:00,  1.14steps/s, loss=0.184]
 32%|███▏      | 32/100 [00:33<00:58,  1.16steps/s, loss=0.184]
 32%|███▏      | 32/100 [00:33<00:58,  1.16steps/s, loss=0.177]
 33%|███▎      | 33/100 [00:34<00:57,  1.16steps/s, loss=0.177]
 33%|███▎      | 33/100 [00:34<00:57,  1.16steps/s, loss=0.17]
 34%|███▍      | 34/100 [00:34<00:56,  1.17steps/s, loss=0.17]
 34%|███▍      | 34/100 [00:35<00:56,  1.17steps/s, loss=0.163]
 35%|███▌      | 35/100 [00:35<00:55,  1.17steps/s, loss=0.163]
 35%|███▌      | 35/100 [00:35<00:55,  1.17steps/s, loss=0.156]
 36%|███▌      | 36/100 [00:36<00:54,  1.17steps/s, loss=0.156]
 36%|███▌      | 36/100 [00:36<00:54,  1.17steps/s, loss=0.15]
 37%|███▋      | 37/100 [00:37<00:59,  1.06steps/s, loss=0.15]
 37%|███▋      | 37/100 [00:37<00:59,  1.06steps/s, loss=0.145]
 38%|███▊      | 38/100 [00:38<00:56,  1.11steps/s, loss=0.145]
 38%|███▊      | 38/100 [00:38<00:56,  1.11steps/s, loss=0.139]
 39%|███▉      | 39/100 [00:39<00:54,  1.13steps/s, loss=0.139]
 39%|███▉      | 39/100 [00:39<00:54,  1.13steps/s, loss=0.134]
 40%|████      | 40/100 [00:40<00:52,  1.15steps/s, loss=0.134]
 40%|████      | 40/100 [00:40<00:52,  1.15steps/s, loss=0.13]
 41%|████      | 41/100 [00:41<00:50,  1.16steps/s, loss=0.13]
 41%|████      | 41/100 [00:41<00:50,  1.16steps/s, loss=0.13]
 42%|████▏     | 42/100 [00:41<00:49,  1.17steps/s, loss=0.13]
 42%|████▏     | 42/100 [00:42<00:49,  1.17steps/s, loss=0.125]
 43%|████▎     | 43/100 [00:42<00:48,  1.17steps/s, loss=0.125]
 43%|████▎     | 43/100 [00:42<00:48,  1.17steps/s, loss=0.12]
 44%|████▍     | 44/100 [00:43<00:47,  1.17steps/s, loss=0.12]
 44%|████▍     | 44/100 [00:43<00:47,  1.17steps/s, loss=0.114]
 45%|████▌     | 45/100 [00:44<00:46,  1.18steps/s, loss=0.114]
 45%|████▌     | 45/100 [00:44<00:46,  1.18steps/s, loss=0.111]
 46%|████▌     | 46/100 [00:45<00:45,  1.18steps/s, loss=0.111]
 46%|████▌     | 46/100 [00:45<00:45,  1.18steps/s, loss=0.109]
 47%|████▋     | 47/100 [00:46<00:49,  1.07steps/s, loss=0.109]
 47%|████▋     | 47/100 [00:46<00:49,  1.07steps/s, loss=0.105]
 48%|████▊     | 48/100 [00:47<00:48,  1.06steps/s, loss=0.105]
 48%|████▊     | 48/100 [00:47<00:48,  1.06steps/s, loss=0.101]
 49%|████▉     | 49/100 [00:48<00:46,  1.09steps/s, loss=0.101]
 49%|████▉     | 49/100 [00:48<00:46,  1.09steps/s, loss=0.0989]
 50%|█████     | 50/100 [00:49<00:46,  1.08steps/s, loss=0.0989]
 50%|█████     | 50/100 [00:49<00:46,  1.08steps/s, loss=0.0948]
 51%|█████     | 51/100 [00:50<00:44,  1.09steps/s, loss=0.0948]
 51%|█████     | 51/100 [00:50<00:44,  1.09steps/s, loss=0.0931]
 52%|█████▏    | 52/100 [00:51<00:44,  1.09steps/s, loss=0.0931]
 52%|█████▏    | 52/100 [00:51<00:44,  1.09steps/s, loss=0.0929]
 53%|█████▎    | 53/100 [00:51<00:42,  1.11steps/s, loss=0.0929]
 53%|█████▎    | 53/100 [00:52<00:42,  1.11steps/s, loss=0.0894]
 54%|█████▍    | 54/100 [00:52<00:40,  1.13steps/s, loss=0.0894]
 54%|█████▍    | 54/100 [00:52<00:40,  1.13steps/s, loss=0.0881]
 55%|█████▌    | 55/100 [00:53<00:39,  1.15steps/s, loss=0.0881]
 55%|█████▌    | 55/100 [00:53<00:39,  1.15steps/s, loss=0.088]
 56%|█████▌    | 56/100 [00:54<00:38,  1.16steps/s, loss=0.088]
 56%|█████▌    | 56/100 [00:54<00:38,  1.16steps/s, loss=0.0842]
 57%|█████▋    | 57/100 [00:55<00:40,  1.06steps/s, loss=0.0842]
 57%|█████▋    | 57/100 [00:55<00:40,  1.06steps/s, loss=0.0791]
 58%|█████▊    | 58/100 [00:56<00:38,  1.10steps/s, loss=0.0791]
 58%|█████▊    | 58/100 [00:56<00:38,  1.10steps/s, loss=0.0758]
 59%|█████▉    | 59/100 [00:57<00:36,  1.12steps/s, loss=0.0758]
 59%|█████▉    | 59/100 [00:57<00:36,  1.12steps/s, loss=0.074]
 60%|██████    | 60/100 [00:58<00:35,  1.14steps/s, loss=0.074]
 60%|██████    | 60/100 [00:58<00:35,  1.14steps/s, loss=0.0735]
 61%|██████    | 61/100 [00:58<00:33,  1.15steps/s, loss=0.0735]
 61%|██████    | 61/100 [00:59<00:33,  1.15steps/s, loss=0.0733]
 62%|██████▏   | 62/100 [00:59<00:33,  1.15steps/s, loss=0.0733]
 62%|██████▏   | 62/100 [00:59<00:33,  1.15steps/s, loss=0.0759]
 63%|██████▎   | 63/100 [01:00<00:31,  1.16steps/s, loss=0.0759]
 63%|██████▎   | 63/100 [01:00<00:31,  1.16steps/s, loss=0.0792]
 64%|██████▍   | 64/100 [01:01<00:31,  1.16steps/s, loss=0.0792]
 64%|██████▍   | 64/100 [01:01<00:31,  1.16steps/s, loss=0.0758]
 65%|██████▌   | 65/100 [01:02<00:30,  1.16steps/s, loss=0.0758]
 65%|██████▌   | 65/100 [01:02<00:30,  1.16steps/s, loss=0.0697]
 66%|██████▌   | 66/100 [01:03<00:29,  1.16steps/s, loss=0.0697]
 66%|██████▌   | 66/100 [01:03<00:29,  1.16steps/s, loss=0.0677]
 67%|██████▋   | 67/100 [01:04<00:28,  1.16steps/s, loss=0.0677]
 67%|██████▋   | 67/100 [01:04<00:28,  1.16steps/s, loss=0.0673]
 68%|██████▊   | 68/100 [01:05<00:31,  1.02steps/s, loss=0.0673]
 68%|██████▊   | 68/100 [01:05<00:31,  1.02steps/s, loss=0.0673]
 69%|██████▉   | 69/100 [01:06<00:29,  1.04steps/s, loss=0.0673]
 69%|██████▉   | 69/100 [01:06<00:29,  1.04steps/s, loss=0.0688]
 70%|███████   | 70/100 [01:07<00:28,  1.06steps/s, loss=0.0688]
 70%|███████   | 70/100 [01:07<00:28,  1.06steps/s, loss=0.0663]
 71%|███████   | 71/100 [01:08<00:26,  1.09steps/s, loss=0.0663]
 71%|███████   | 71/100 [01:08<00:26,  1.09steps/s, loss=0.0603]
 72%|███████▏  | 72/100 [01:08<00:25,  1.11steps/s, loss=0.0603]
 72%|███████▏  | 72/100 [01:09<00:25,  1.11steps/s, loss=0.0584]
 73%|███████▎  | 73/100 [01:09<00:24,  1.12steps/s, loss=0.0584]
 73%|███████▎  | 73/100 [01:09<00:24,  1.12steps/s, loss=0.0584]
 74%|███████▍  | 74/100 [01:10<00:22,  1.14steps/s, loss=0.0584]
 74%|███████▍  | 74/100 [01:10<00:22,  1.14steps/s, loss=0.06]
 75%|███████▌  | 75/100 [01:11<00:21,  1.15steps/s, loss=0.06]
 75%|███████▌  | 75/100 [01:11<00:21,  1.15steps/s, loss=0.0624]
 76%|███████▌  | 76/100 [01:12<00:20,  1.16steps/s, loss=0.0624]
 76%|███████▌  | 76/100 [01:12<00:20,  1.16steps/s, loss=0.06]
 77%|███████▋  | 77/100 [01:13<00:19,  1.17steps/s, loss=0.06]
 77%|███████▋  | 77/100 [01:13<00:19,  1.17steps/s, loss=0.0578]
 78%|███████▊  | 78/100 [01:14<00:20,  1.07steps/s, loss=0.0578]
 78%|███████▊  | 78/100 [01:14<00:20,  1.07steps/s, loss=0.055]
 79%|███████▉  | 79/100 [01:15<00:19,  1.11steps/s, loss=0.055]
 79%|███████▉  | 79/100 [01:15<00:19,  1.11steps/s, loss=0.054]
 80%|████████  | 80/100 [01:15<00:17,  1.13steps/s, loss=0.054]
 80%|████████  | 80/100 [01:16<00:17,  1.13steps/s, loss=0.0557]
 81%|████████  | 81/100 [01:16<00:16,  1.15steps/s, loss=0.0557]
 81%|████████  | 81/100 [01:16<00:16,  1.15steps/s, loss=0.0583]
 82%|████████▏ | 82/100 [01:17<00:15,  1.16steps/s, loss=0.0583]
 82%|████████▏ | 82/100 [01:17<00:15,  1.16steps/s, loss=0.0556]
 83%|████████▎ | 83/100 [01:18<00:14,  1.17steps/s, loss=0.0556]
 83%|████████▎ | 83/100 [01:18<00:14,  1.17steps/s, loss=0.0517]
 84%|████████▍ | 84/100 [01:19<00:13,  1.18steps/s, loss=0.0517]
 84%|████████▍ | 84/100 [01:19<00:13,  1.18steps/s, loss=0.0516]
 85%|████████▌ | 85/100 [01:20<00:12,  1.18steps/s, loss=0.0516]
 85%|████████▌ | 85/100 [01:20<00:12,  1.18steps/s, loss=0.0549]
 86%|████████▌ | 86/100 [01:20<00:11,  1.18steps/s, loss=0.0549]
 86%|████████▌ | 86/100 [01:21<00:11,  1.18steps/s, loss=0.0558]
 87%|████████▋ | 87/100 [01:21<00:10,  1.19steps/s, loss=0.0558]
 87%|████████▋ | 87/100 [01:21<00:10,  1.19steps/s, loss=0.0531]
 88%|████████▊ | 88/100 [01:22<00:11,  1.08steps/s, loss=0.0531]
 88%|████████▊ | 88/100 [01:23<00:11,  1.08steps/s, loss=0.0559]
 89%|████████▉ | 89/100 [01:23<00:10,  1.10steps/s, loss=0.0559]
 89%|████████▉ | 89/100 [01:23<00:10,  1.10steps/s, loss=0.0539]
 90%|█████████ | 90/100 [01:24<00:08,  1.12steps/s, loss=0.0539]
 90%|█████████ | 90/100 [01:24<00:08,  1.12steps/s, loss=0.0518]
 91%|█████████ | 91/100 [01:25<00:07,  1.14steps/s, loss=0.0518]
 91%|█████████ | 91/100 [01:25<00:07,  1.14steps/s, loss=0.0509]
 92%|█████████▏| 92/100 [01:26<00:06,  1.15steps/s, loss=0.0509]
 92%|█████████▏| 92/100 [01:26<00:06,  1.15steps/s, loss=0.0505]
 93%|█████████▎| 93/100 [01:27<00:06,  1.14steps/s, loss=0.0505]
 93%|█████████▎| 93/100 [01:27<00:06,  1.14steps/s, loss=0.0501]
 94%|█████████▍| 94/100 [01:28<00:05,  1.16steps/s, loss=0.0501]
 94%|█████████▍| 94/100 [01:28<00:05,  1.16steps/s, loss=0.0477]
 95%|█████████▌| 95/100 [01:28<00:04,  1.17steps/s, loss=0.0477]
 95%|█████████▌| 95/100 [01:29<00:04,  1.17steps/s, loss=0.0483]
 96%|█████████▌| 96/100 [01:29<00:03,  1.18steps/s, loss=0.0483]
 96%|█████████▌| 96/100 [01:29<00:03,  1.18steps/s, loss=0.0508]
 97%|█████████▋| 97/100 [01:30<00:02,  1.18steps/s, loss=0.0508]
 97%|█████████▋| 97/100 [01:30<00:02,  1.18steps/s, loss=0.0516]
 98%|█████████▊| 98/100 [01:31<00:01,  1.18steps/s, loss=0.0516]
 98%|█████████▊| 98/100 [01:31<00:01,  1.18steps/s, loss=0.0501]
 99%|█████████▉| 99/100 [01:32<00:00,  1.08steps/s, loss=0.0501]
 99%|█████████▉| 99/100 [01:32<00:00,  1.08steps/s, loss=0.0466]
100%|██████████| 100/100 [01:33<00:00,  1.11steps/s, loss=0.0466]
100%|██████████| 100/100 [01:33<00:00,  1.07steps/s, loss=0.0466]
example learn_samples

Reconstruction from partially trained U-Net model

model.eval()
new_recon = model(kspace_mri_2D)
fig, axs = plt.subplots(2, 2, figsize=(10, 10))
plot_state(axs, mri_2D, init_traj, new_recon, losses)
plt.show()
MR Image, Trajectory, Reconstruction, Loss

References#

[fastmri] (1,2)

O. Ronneberger, P. Fischer, and Thomas Brox. U-net: Convolutional networks for biomedical image segmentation. In International Conference on Medical image computing and computer-assisted intervention, pages 234–241. Springer, 2015. facebookresearch/fastMRI

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

Gallery generated by Sphinx-Gallery