ei_demo_cs_usps.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "name": "ei_demo_cs_usps.ipynb",
      "provenance": [],
      "collapsed_sections": [],
      "authorship_tag": "ABX9TyNl9N136Zp6YCs/FNF0KeLD",
      "include_colab_link": true
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "language_info": {
      "name": "python"
    },
    "accelerator": "GPU"
  },
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "view-in-github",
        "colab_type": "text"
      },
      "source": [
        "<a href=\"https://colab.research.google.com/github/edongdongchen/EI/blob/main/ei_demo_cs_usps.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "OjnZAd-LuM_B"
      },
      "source": [
        "import numpy as np\n",
        "import torch\n",
        "import torch.nn as nn\n",
        "import torchvision.transforms as transforms\n",
        "from torchvision.datasets import USPS\n",
        "from torch.utils.data import DataLoader\n",
        "\n",
        "import os\n",
        "from datetime import datetime\n",
        "import matplotlib.pyplot as plt"
      ],
      "execution_count": 1,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "razUOZiJy_ip"
      },
      "source": [
        "# define an inverse problem e.g. Compressed Sensing (CS)\n",
        "# where the forward operator A is a random projection matrix\n",
        "class CS():\n",
        "  def __init__(self, d, D, img_shape, dtype=torch.float, device='cuda:0'):\n",
        "    self.img_shape = img_shape\n",
        "    fname = './dataset/forw_cs_{}x{}.pt'.format(d, D)\n",
        "    if os.path.exists(fname):\n",
        "      A, A_dagger = torch.load(fname)\n",
        "    else:\n",
        "      A = np.random.randn(d, D) / np.sqrt(d)\n",
        "      A_dagger = np.linalg.pinv(A)\n",
        "      torch.save([A, A_dagger], fname)\n",
        "      print('CS matrix has been created and saved at {}'.format(fname))\n",
        "    self._A = torch.from_numpy(A).type(dtype).to(device)\n",
        "    self._A_dagger = torch.from_numpy(A_dagger).type(dtype).to(device)\n",
        "\n",
        "  def A(self, x):\n",
        "    y = torch.einsum('in, mn->im', x.reshape(x.shape[0], -1), self._A)\n",
        "    return y\n",
        "\n",
        "  def A_dagger(self, y):\n",
        "    N, C, H, W = y.shape[0], self.img_shape[0], self.img_shape[1], self.img_shape[2]\n",
        "    x = torch.einsum('im, nm->in', y, self._A_dagger)\n",
        "    x = x.reshape(N, C, H, W)\n",
        "    return x"
      ],
      "execution_count": 2,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "n9NGveNuTFXo"
      },
      "source": [
        "Plot the compressed measurement $y$ and its Groundtruth $x$. "
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 231
        },
        "id": "WPtgnuYMw-tj",
        "outputId": "0d702ab2-826d-490f-f14e-a9ca2761d618"
      },
      "source": [
        "dtype = torch.float\n",
        "device = 'cuda:0'\n",
        "batch_size = 1\n",
        "\n",
        "# test set\n",
        "test_dataset = USPS(root='./dataset/', train=False, download=True, transform=transforms.ToTensor())\n",
        "test_dataloader=DataLoader(test_dataset,batch_size=1,shuffle=False)\n",
        "\n",
        "# define a CS (4x) operator\n",
        "physics = CS(64, 256, [1,16,16], dtype, device)\n",
        "# ground-truh\n",
        "x = test_dataset[0][0].type(dtype).to(device) #torch.Size([1, 1, 16, 16])\n",
        "# measurement\n",
        "y = physics.A(x)\n",
        "\n",
        "print(f'x (GT): {x.shape}\\ty: {y.shape}')\n",
        "\n",
        "def torch2img(x):\n",
        "  return x.squeeze().detach().cpu().numpy()\n",
        "  \n",
        "def cal_psnr(a, b):\n",
        "  alpha = np.sqrt(a.shape[-1]*a.shape[-2])\n",
        "  psnr = 20*torch.log10(alpha*torch.norm(b, float('inf'))/torch.norm(b-a, 2))\n",
        "  return psnr.detach().cpu().numpy()\n",
        "\n",
        "ATy = physics.A_dagger(y)\n",
        "# have a look at the cs measurement y, A^+y, and the ground-truth x.\n",
        "plt.subplot(1,3,1)\n",
        "plt.imshow(torch2img(y.view(1,1,4,16)), cmap='gray'), plt.title('y'), plt.axis('off')\n",
        "plt.subplot(1,3,2)\n",
        "plt.imshow(torch2img(ATy), cmap='gray'), plt.title(r'$A^+y$ ({:.2f} dB)'.format(cal_psnr(ATy, x))), plt.axis('off')\n",
        "plt.subplot(1,3,3)\n",
        "plt.imshow(torch2img(x), cmap='gray'), plt.title('x (GT)'), plt.axis('off')\n",
        "plt.show()"
      ],
      "execution_count": 3,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Downloading https://www.csie.ntu.edu.tw/~cjlin/libsvmtools/datasets/multiclass/usps.t.bz2 to ./dataset/usps.t.bz2\n"
          ]
        },
        {
          "output_type": "display_data",
          "data": {
            "application/vnd.jupyter.widget-view+json": {
              "model_id": "67435fc9fa7a409c9eae1fb077dd0671",
              "version_minor": 0,
              "version_major": 2
            },
            "text/plain": [
              "  0%|          | 0/1831726 [00:00<?, ?it/s]"
            ]
          },
          "metadata": {}
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "CS matrix has been created and saved at ./dataset/forw_cs_64x256.pt\n",
            "x (GT): torch.Size([1, 16, 16])\ty: torch.Size([1, 64])\n"
          ]
        },
        {
          "output_type": "display_data",
          "data": {
            "image/png": "iVBORw0KGgoAAAANSUhEUgAAAV0AAACDCAYAAAAuy8hcAAAABHNCSVQICAgIfAhkiAAAAAlwSFlzAAALEgAACxIB0t1+/AAAADh0RVh0U29mdHdhcmUAbWF0cGxvdGxpYiB2ZXJzaW9uMy4yLjIsIGh0dHA6Ly9tYXRwbG90bGliLm9yZy+WH4yJAAASYElEQVR4nO3deXBVZZoG8OdlkR0iYRGiiEE2hQg4qCjKMnYaghSCjuI02DhQVjsy49Y6ypTVTmszKDM0NoUNQ+kA0sA4YHBpelAGkJLCZqtW1g6oAZRFFAJJCIvwzh/n4FzS73sxEL4bb55fFVXkuWf5bu697zk557vfJ6oKIiIKo0aqG0BEVJ2w6BIRBcSiS0QUEIsuEVFALLpERAGx6BIRBcSiS0QUEIsuVQoRmZnqNhD9ELDoUpUjIv8qIo+muh3lichMEXmhAsuvEZFrL2ab6IeHRZdcIlJLRLaKyPaA+2wO4H4A0+OfO4vIMhE5LCI7RGToOdZvKiL5IlIqIjtF5G+d5dqLyDERmXMBbS0UkTIRKRGRQyLyexG5ImGRfwPwy/PdfnVQmQdYEakjItvi91CVxaJLyfwMQAsA2SLSoPyDItJGRFaIyAoAA878X0QuuYB9jgKwWFXLRKQWgLcAvAugKYAHAcwRkQ5J1p8K4ASAlgB+AuC3ztnmVABrL6CdZwxW1YYAWgHYD2BKwmNvA+gnIpdVwn7STvkDbJw1EpFJ8QGtVER2icgCEblRRMbFB7iS+IB5KuHnzap6HMBrAJ5O1XP6Plh0ySQiTQD8AsBYAKcAdCm/jKruUtW+qtoXwP+c+b+qnkjYTsP4w9EqIesiIntFpJGx64EAPoj/3wlAawC/VtVTqroMwCoAI502NwBwF4BnVbVEVT9EVPhGlltuOIAiAP97jt9BdxHZICLFIvJfAOp6y6rqMQALAFxTLlsP4MfJ9lONjUJ8gAWiM1UAywB0BXAHgMYAOgOYD2Cgqo5X1YbxQe5nAFaf+VlVzxxY5wL4abytKolFlzzPAvhMVecBKACQcz4bUdUSANsA9EiIJwAYr6rFxipdAfw5ySYFxgEg1gHAt6pakJB9DOC7M10RaYzoT/7Hk7U7PltfBOB1RGfZ/42ooHvL1wdwL4CPyj20FcB1yfaVLkSknYgcFJEe8c+tReSAiPR1Vkk8wALRwfFyAHeq6qb4QFuqqgtU9bnv0wZV/QLAIQA3nfcTuchYdOkviEg2gH8A8E9xtBnnKLqqOirJw2sRF10RuQ3R2eB0Z9kMAGeK8Z8BfAXgSRGpLSK5APoAqO+s2xDAkXLZYQCJZ9TPA3g1/nAmcxOA2gAmq+pJVV0A+3LEIhEpivfzIwATyz1eHD+ntKeqnyJ6z8yJD0L/CWCWqq5wVil/gL0dwBJVLb3AplTpAx2LLlleArA04cOyGRf2Jv6u6MbbfjbxEkQ5hxAXSVU9CeBOAIMA7APwBIA3AHgFswTRn6SJGiMu4iLSDdEH+9ffo82tAXypZ499utNY7k5VzUB06WEsgA/KXcNthOhSRrWgqjMA7ADwR0TXuf85yeKJB1gAaIbodQYQvV4iUiQiR0Qk2V8/5VXpAx2LLp1FRHoj+jO6t4jsE5Ezxa7rBWx2LYAeInIXouI0N8mynyC6TAAAUNVPVLWPqmaq6o8BZANY46xbAKCWiLRPyK5DdNAAgL4A2gLYFT+vnwO4S0Q2GNvaCyBLRCQha+M1Ov5T+E1E1797JzzUGdEljupkBqJLQFPim1ue7w6wsW8QFWoAgKr+KT6gDQNQkWu0VfpAx6JL34kLzCQA0wB0BNAt/pcLIENE3KJzDh8DuAzAvwN4ptzZY3mLEV1CONOmHBGpKyL1ReTniD6UM60V4z9L3wTwSxFpICK3ABiC6LosAPwHgHYJz2sagN/DvtG1GsC3AP4xvrQxDMANXqMlMgTApYj+vIWI1AVwPYD3kzzftCIiDQFMBvAqgOdEpGmSxc86wCK6sZlr9ZSpoCp9oGPRpUQ/QdTV6klV3XfmH6Izy2Kc/8204wA2AihU1T+cY/HZAPJEpF7880hEZ51fAfhrAD9KPHsSkT+IyLiE9f8eQL14+XkAHlLVzXE7jpZ7XiUAjqnqAaPNJxCdYY0CcBDRTbI3jfa+IyIliK4l/wrAT8/sD8BgACtUdc85nnM6eRnAOlUdg+iANi3JsmcdYBG99nsB5Mc9XGrGB66/+r47F5EsRDc+y9/QrDKE0/XQxRb3BNgB4B5VPeeHQUTGA/hKVSdf9MZdRCLyRwCjVXVTqtsSQnym/wqArqp6MD7r/ROAX6jq74zlm8WPt0/oNtYEwL8AGIroGu/XANYBeFFV1ySsOwrAGFXtXW6bTwJopapJe6ekEosuXXQi8isA2ap6X6rbQlVLZR5g4765HwO4TVW/uuDGXSQsunTRxP01lyO6djdUVb9OcZOIUo5Fl4goIN5IIyIKiEWXiCigWqluAIX3+OOPm9eU8vPzzeVbtGjhbiszM9PMN2ywvm8A9O/fv0L7qFHDPi/o3bu3mQNATo7ds23VqlVmvmPHDjN/9dVXzfymmyr+tf66de2xcubOnSvmA+dJRCrtemHNmjXNPCPD/rJXu3btzPyaa64x806dOpl5VlaW26ZGjawxkoDCwkIz/+CDD8x83bp1Zr5v3z4zB4CTJ0+6j1lU1XxteaZLRBQQiy4RUUAsukREAbHoEhEFxKJLRBQQey9UQ0eOlB/nO9K5c2czb9Kkibut3bt3V2idKVOmmPnGjRvN/MCBvxiLBgAwY8YMt01Dhgwx882bN5u517NgyZIlZv7WW2+5+/Z6c8yePdtdJ5W8HgoA0Ly5Pb+j13Nk8ODBZt6xY8cK7buivQQA//fepo09MJ7XA+O9995z9+H1bDh16tQ5Wnc2nukSEQXEoktEFBCLLhFRQCy6REQBsegSEQXE3gvV0Lfffmvmx44dM/PatWu72/K+V3/rrbea+fLly818zRp7rsmrr77azB977DG3TVdeeaWZf/7552bevn17M//4Y3uaraIif87DefPmmXleXp67Tghnz6/5/5o29acw69Onj5nfe++9Zt62bVsz//LLL83cG59j+/btbpvKysrM/KqrrjJz7zkMGDDAzL/+2h/yeeXKlWZ++PBhdx0Lz3SJiAJi0SUiCohFl4goIBZdIqKAWHSJiAJi0SUiCohdxqohr0uO1zXslltucbflTbOzadMmM3/ooYfMvKCgwMznz59v5t4gKgBw6NAhM/emh/EG7fF+T94gKgDw+uuvm/kdd9xh5hMnTnS3VZnq1atn5tddd527zrBhw8z88ssvN/Nly5aZ+dKlS81827ZtZp6sC5Y3fdPevXvN3HutvEF7+vbt6+5769atZl5cXOyuY+GZLhFRQCy6REQBsegSEQXEoktEFBCLLhFRQGnZe0FEngRwk6relZD9BoCq6iOpa1nV0KtXLzNv0KCBmT/66KPutrxeB97AITfccIOZewOW9OzZ08zr16/vtsnrOXHJJZeYeU5Ojpl7d++9AYMA4JVXXjHz2267zV2nMnkD23i9TLyBiQB/sKF169aZeX5+vpl7AwcdPXrUzL0eCgCQlZVl5l27djVz7zm0bt3azLOzs919N27c2My937knXc905wAYICIZACAitQAMB1A1J6oiomojLYuuqu4FsBLA38TRAABfq+r61LWKiChNi25sFoAR8f9HALB7rRMRBZTORXcRgBwR6QLgDgC/S3F7iIjSt+iq6jEACwDMBbBGVXeluElEROnZeyHBLABjAPxdqhvyQ7B48WIzr1u3rruONwWO13uhVatWZj548GAz37Jli5kvWbLEbdMjj9gdVLzeGV7uPe/Zs/37sZmZmWZep04dd53K5N35z8jIMHNvuiUAOH36tJl/9NFHZu6NpVDRXgpeWwGgR48eZu5NIdS9e3czP3LkiJnv2uWfm5WWlpq5qrrrWNL2TDe2C0AZgIWpbggREZDGRVdEagB4HMB8VbUPa0REgaXl5QURaQBgP4CdiLqLERFVCWlZdFW1FEDDVLeDiKi8tL28QERUFSU9061fv755W27s2LHm8vfff7+7LW/GgNGjR5u5d1fcG6W9X79+Zu59Dx8Axo0bZ+YtW7Y085deesnMi4qKzNz7rjYAfPHFFxXKp02bVrEveCdRWFho5t73zmvWrOlua+fOnWael5dn5t5d5vfff9/MvXEfXnjhBbdNp06dMnPvPeWNB+Etn2wcBW+WhGR35EOoVcv+qHvjUQB+r4N9+/aZufd7b9q0qZk3a9bMzK+99lq3TV4vF28GDK9N3iwXb7/9trvvPXv2mLnXy8PDM10iooBYdImIAmLRJSIKiEWXiCggFl0iooCS9l4YMWKEmXtzzHvLA8DQoUPNfM2aNWbufbd+5syZZv7yyy+bebJxA7xR5V988UUz93pCbN261czfffddd9+XXXaZmW/fvt1dp7J4Mx9434WvXbu2uy2vN4L3HfZFixaZ+fPPP2/mXg+TSZMmuW3yxn3o0qWLmT/11FNmXtFZGACgU6dOZu71GKls3jgAJSUlZu71lgGAhg3tru7ec/TeJ954G964D954CYDfs8EbP2PlypVm7s1ysWHDBnff3tgLFcUzXSKigFh0iYgCYtElIgqIRZeIKCAWXSKigCTZqOfdunUzH/TGP7jxxhvdbXl3+L1teWMyLFxoj0fufR/c6zUBALfffruZd+zY0cy9O+nendMxY8a4+/bGGvC+15+bm1tpYy+IiPm6jhw50lzeu/MN+M89KyvLzL3eDl6vhm7dupl5QUGB26bPPvvMzL2xJebPn2/mDz/8sJl7MyQAQIcOHcx8xowZZr527dpKe10BoEaNGuZr27x5c3P5QYMGudsaNmyYmXtjKXhjdHizZnjjQXjbB4DWrVubuff5nzp1qpnPmTPHzHfv3u3u2xvHwaOq5mvLM10iooBYdImIAmLRJSIKiEWXiCggFl0iooBYdImIAko64I03xY43SMaqVavcbU2cONHMvSk7JkyYYObeFCretCPt2rVz2+RNR+INerF//34z9wZSKSsrc/d97NgxM1++fLmZ5+bmutuqqKefftrMva5QyaZh2rFjh5l7XQS9AUsmT55s5t6gQePHj3fbdPLkSTP3uhX17dvXzL0pXbwuhQDwzDPPmPndd9/trlOZvC6ghw4dMnNvQBjAn4amV69eZu4NbOO9Ht7ARN52AODmm282c6+r18aNG838wIEDZl7RbmHng2e6REQBsegSEQXEoktEFBCLLhFRQCy6REQBJR3whtLTAw88YL7oXg8Qr3cGABw8eNDMN2/eXKHlr7jiCjMfMmSImb/xxhtum5544gkznz59upl7z88bVMebxgfwB5Z57bXXzLywsLBSB7zxBjPyJJuKKSMjw8wzMzPN3Bv8yONt57777nPX6devn5mvXr3azL1eU5988omZe1NZnQ8OeENEVAWw6BIRBcSiS0QUEIsuEVFALLpERAElHXuB0pN317hx48Zm/uGHH7rb6t69u5l70/V437f3viPvjZ3h7RcA1q9fb+bedDLe9+1LS0vNPFnvhSZNmpj5Pffc466TSt64CID/Wn3zzTdm7v1evGl5rr/+ejNv0aKF2yZvPAhvDBCvt0yIMRY8PNMlIgqIRZeIKCAWXSKigFh0iYgCYtElIgqIvReqoS1btph5Xl6emRcXF7vbOnLkiJmvWLHCzIcPH27mhw8fNnOvB0F+fr7bJm9mg+3bt5v54sWLzbxt27ZmXqdOHXff3owgycavqKq8cVkqOl6LN76D19MjWe8F7/3gvbbe+yqVY87wTJeIKCAWXSKigFh0iYgCYtElIgqIRZeIKCD2XqiG9u7da+Zjx441808//dTd1jvvvGPm2dnZZu71hFi7dq2Zez0qks0ukJuba+YTJkww8zZt2ph5SUmJmTdr1szdt3cXvWXLlu466c7r7eHNTJGsd4j3/tmzZ4+ZHz169BytC49nukREAbHoEhEFxKJLRBQQiy4RUUAsukREAbHoEhEFxC5j1ZDXJee5554zc296FgAYNGiQmZ84ccLMFy1aZOY5OTlm7nXB8gaWAfzuQ0OHDjXz1atXm/nSpUvNvH///u6+vemFvCmERo8e7W4rXXjT+HiDziQbYMmbfsfr3pfKaXk8PNMlIgqIRZeIKCAWXSKigFh0iYgCYtElIgqIvReqoQcffNDMZ82aZeaXXnqpu62CggIzP378uJl36NDBzL1BZ1atWmXmPXv2dNvkbWvhwoXuOhZvqpdkg9d4U/8MHDiwQvtOJ2VlZWbu9UxZtmyZuy2vx0NRUZGZnz59+hytC49nukREAbHoEhEFxKJLRBQQiy4RUUAsukREAYl3N5CIiCofz3SJiAJi0SUiCohFl4goIBZdIqKAWHSJiAJi0SUiCuj/AAAlfl/6UVV3AAAAAElFTkSuQmCC\n",
            "text/plain": [
              "<Figure size 432x288 with 3 Axes>"
            ]
          },
          "metadata": {
            "needs_background": "light"
          }
        }
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "AE-rRMOm93Lj"
      },
      "source": [
        "\n",
        "\n",
        "---\n",
        "To learn the reconstruction function $f:y\\rightarrow x$ from only the measurements set $\\{y\\}$, we assume the signal set $\\mathcal{X}$ is invariant to a groups of transformations $T_g$ related to a group 𝒢, where $T_gT_g^{\\top}=I$. For example:\n",
        "    \n",
        "    - natural images are shift invariant.\n",
        "    - in CT/MRI data, organs can be imaged at different angles making the problem invariant to rotation.\n",
        "\n",
        "Key observations: \n",
        "- Invariance provides access to implicit operators $A_g=AT_g$ with potentially different range spaces: $Ax=AT_gT_g^{\\top}x=A_g\\tilde{x}$  where $A_g=AT_g$ and $\\tilde{x}=T_g^{\\top}x$.\n",
        "\n",
        "- The composition $f\\circ A$ is equivariant to the group of transformations ${T_g}$:  $f(AT_gx)=T_gf(Ax)$\n",
        "\n",
        "\n",
        "# **EI:** All you need is to:\n",
        "  * Define:\n",
        "1. define a transformation group $\\{T_g\\}$ based on the certain invariances to the signal set.\n",
        "2. define a neural reconstruction function $f_\\theta: y\\rightarrow x$, *e.g.* $f_\\theta=G_\\theta\\circ A^{\\dagger}$ and $G_\\theta$ could be a UNet-like encoder-decoder net.\n",
        "\n",
        "  * Calculate:\n",
        "1. calculate $x^{(1)} = f_\\theta(y)$ as the estimation of $x$.\n",
        "2. calculate $x^{(2)} = T_g(x^{(1)})$ by transforming $x^{(1)}$.\n",
        "2. calculate $x^{(3)} = f_\\theta(Ax^{(2)})$ by reconstructing $x^{(2)}$ from its measurement $Ax^{(2)}$.\n",
        "\n",
        "\n",
        "  * Finaly, learn the reconstruction function $f_\\theta$ by solving: \n",
        "$\\arg\\min\\limits_\\theta\\mathbb{E}_{y,g}\\{\\mathcal{L}(Ax^{(1)},y) +\\lambda\\mathcal{L}(x^{(2)},x^{(3)})\\}$."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "bqiygCZ0HcQs"
      },
      "source": [
        "# define a transformation group (random shift)\n",
        "import random\n",
        "class Shift():\n",
        "  def __init__(self, n_trans):\n",
        "    self.n_trans = n_trans\n",
        "  def apply(self, x):\n",
        "    H, W = x.shape[-2], x.shape[-1]\n",
        "    assert self.n_trans <= H - 1 and self.n_trans <= W - 1, 'n_shifts should less than {}'.format(H-1)\n",
        "\n",
        "    shifts_row = random.sample(list(np.concatenate([-1*np.arange(1, H), np.arange(1, H)])), self.n_trans)\n",
        "    shifts_col = random.sample(list(np.concatenate([-1*np.arange(1, W), np.arange(1, W)])), self.n_trans)\n",
        "\n",
        "    x = torch.cat([x if self.n_trans == 0 else torch.roll(x, shifts=[sx, sy], dims=[-2,-1]).type_as(x)\n",
        "              for sx, sy in zip(shifts_row, shifts_col)], dim=0)\n",
        "    return x"
      ],
      "execution_count": 4,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "tUX5jch5UeBf"
      },
      "source": [
        "# define residual unet\n",
        "class UNet(nn.Module):\n",
        "  def __init__(self, in_channels=1, out_channels=1, compact=4, residual=True, circular_padding=True, cat=True):\n",
        "    super(UNet, self).__init__()\n",
        "    self.name = 'unet'\n",
        "    self.residual = residual\n",
        "    self.cat = cat\n",
        "\n",
        "    self.Maxpool = nn.MaxPool2d(kernel_size=2, stride=2)\n",
        "\n",
        "    self.Conv1 = conv_block(ch_in=in_channels, ch_out=64, circular_padding=circular_padding)\n",
        "    self.Conv2 = conv_block(ch_in=64, ch_out=128)\n",
        "    self.Conv3 = conv_block(ch_in=128, ch_out=256)\n",
        "    self.Conv4 = conv_block(ch_in=256, ch_out=512)\n",
        "\n",
        "    self.Up4 = up_conv(ch_in=512, ch_out=256)\n",
        "    self.Up_conv4 = conv_block(ch_in=512, ch_out=256)\n",
        "\n",
        "    self.Up3 = up_conv(ch_in=256, ch_out=128)\n",
        "    self.Up_conv3 = conv_block(ch_in=256, ch_out=128)\n",
        "\n",
        "    self.Up2 = up_conv(ch_in=128, ch_out=64)\n",
        "    self.Up_conv2 = conv_block(ch_in=128, ch_out=64)\n",
        "\n",
        "    self.Conv_1x1 = nn.Conv2d(in_channels=64, out_channels=out_channels, kernel_size=1, stride=1, padding=0)\n",
        "  \n",
        "  def forward(self, x):\n",
        "    # encoding path\n",
        "    cat_dim = 1\n",
        "    input = x\n",
        "    x1 = self.Conv1(input)\n",
        "\n",
        "    x2 = self.Maxpool(x1)\n",
        "    x2 = self.Conv2(x2)\n",
        "\n",
        "    x3 = self.Maxpool(x2)\n",
        "    x3 = self.Conv3(x3)\n",
        "\n",
        "    d3 = self.Up3(x3)\n",
        "    if self.cat:\n",
        "        d3 = torch.cat((x2, d3), dim=cat_dim)\n",
        "        d3 = self.Up_conv3(d3)\n",
        "\n",
        "    d2 = self.Up2(d3)\n",
        "    if self.cat:\n",
        "        d2 = torch.cat((x1, d2), dim=cat_dim)\n",
        "        d2 = self.Up_conv2(d2)\n",
        "\n",
        "    d1 = self.Conv_1x1(d2)\n",
        "\n",
        "    out = d1+x if self.residual else d1\n",
        "    return out\n",
        "\n",
        "class conv_block(nn.Module):\n",
        "  def __init__(self, ch_in, ch_out, circular_padding=False):\n",
        "    super(conv_block, self).__init__()\n",
        "    self.conv = nn.Sequential(\n",
        "        nn.Conv2d(ch_in, ch_out, kernel_size=3, stride=1, padding=1, bias=True,\n",
        "                  padding_mode='circular' if circular_padding else 'zeros'),\n",
        "        nn.BatchNorm2d(ch_out),\n",
        "        nn.ReLU(inplace=True),\n",
        "        nn.Conv2d(ch_out, ch_out, kernel_size=3, stride=1, padding=1, bias=True),\n",
        "        nn.BatchNorm2d(ch_out),\n",
        "        nn.ReLU(inplace=True)\n",
        "    )\n",
        "  def forward(self, x):\n",
        "    x = self.conv(x)\n",
        "    return x\n",
        "\n",
        "class up_conv(nn.Module):\n",
        "  def __init__(self, ch_in, ch_out):\n",
        "    super(up_conv, self).__init__()\n",
        "    self.up = nn.Sequential(\n",
        "        nn.Upsample(scale_factor=2),\n",
        "        nn.Conv2d(ch_in, ch_out, kernel_size=3, stride=1, padding=1, bias=True),\n",
        "        nn.BatchNorm2d(ch_out),\n",
        "        nn.ReLU(inplace=True)\n",
        "    )\n",
        "\n",
        "  def forward(self, x):\n",
        "    x = self.up(x)\n",
        "    return x"
      ],
      "execution_count": 5,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "sMxC3CHeOwZO"
      },
      "source": [
        "Define the Neural Reconstruction Funtion $f_\\theta=G_\\theta\\circ A^\\dagger$:"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "MkhUYkdoGngt"
      },
      "source": [
        "G = UNet().to(device)\n",
        "f = lambda y: G(physics.A_dagger(y))"
      ],
      "execution_count": 6,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "fbNWZqhGYYfI"
      },
      "source": [
        "Define the transformation group $\\{T_g\\}$:\n"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "UGqs8d1dYfNY"
      },
      "source": [
        "T = Shift(n_trans=2) # random shift: return n_trans=2 randomly shifted versions of the input"
      ],
      "execution_count": 7,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Y9IAfN1lPCNa"
      },
      "source": [
        "**Start Training:**"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "AcDi0thJYq71"
      },
      "source": [
        "# Equivariant Imaging\n",
        "# note: the below hyper-parameters value is just for an example and may not be the best setting.\n",
        "lr = 5e-4\n",
        "wd = 1e-8\n",
        "batch_size=8\n",
        "epochs = 300\n",
        "ckp_interval = 100\n",
        "alpha = 10 # equivariance strength\n",
        "\n",
        "\n",
        "dataloader=DataLoader(USPS(root='./dataset/',train=True,download=True, transform=transforms.ToTensor()),batch_size=batch_size,shuffle=True)\n",
        "now = lambda: datetime.now().strftime('%y-%m-%d-%H:%M:%S')\n",
        "\n",
        "def train(alpha=alpha):\n",
        "  optimizer = torch.optim.Adam(G.parameters(), lr=lr, weight_decay=wd)\n",
        "  mse = torch.nn.MSELoss().to(device)\n",
        "\n",
        "  save_path = './ckp/{}_ei_cs4x_alpha_{}'.format(now(), alpha)\n",
        "  os.makedirs(save_path, exist_ok=True)\n",
        "\n",
        "  for epoch in range(epochs):\n",
        "    loss_mc_epoch, loss_ei_epoch, loss_epoch, psnr_epoch = [], [], [], []\n",
        "    for i, x in enumerate(dataloader):\n",
        "      x = x[0].type(dtype).to(device) #torch.Size([4, 1, 16, 16])\n",
        "      y = physics.A(x) #generate the measurements for training\n",
        "      \n",
        "      # only 4 lines of code for training EI\n",
        "      # =========== EI start ============\n",
        "      x1 = f(y) # reconstruction of the groundtruth x\n",
        "      x2 = T.apply(x1) # transform\n",
        "      x3 = f(physics.A(x2)) # reconstruct the transformed data\n",
        "      loss = mse(physics.A(x1), y) + alpha * mse(x3, x2)\n",
        "      # =========== EI end ============\n",
        "\n",
        "      loss_epoch.append(loss.item())\n",
        "      psnr_epoch.append(cal_psnr(x1, x)) # \n",
        "\n",
        "      optimizer.zero_grad()\n",
        "      loss.backward()\n",
        "      optimizer.step()\n",
        "    print('{}\\t[{}/{}]\\tloss={:.4e}\\tpsnr={:.3f}'.format(now(), epoch+1,epochs,np.mean(loss_epoch), np.mean(psnr_epoch)))\n",
        "    \n",
        "    # save the trained model\n",
        "    if epoch % ckp_interval == 0 or epoch + 1 == epochs:\n",
        "      state = {'epoch': epoch, 'state_dict': G.state_dict(), 'optimizer': optimizer.state_dict()} \n",
        "      torch.save(state, os.path.join(save_path, 'ckp_{}.pth.tar'.format(epoch)))\n",
        "\n",
        "train()\n"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "92C0NFvJS2-s"
      },
      "source": [
        "**Test** a trained EI model."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "yTp6vmmEBFXI"
      },
      "source": [
        "# Test\n",
        "def test(net, ckp, ATy):\n",
        "  checkpoint = torch.load(ckp, map_location=device)\n",
        "  net.load_state_dict(checkpoint['state_dict'])\n",
        "  net.to(device).eval()\n",
        "  return net(ATy)\n",
        "\n",
        "# load the trained EI model\n",
        "ckp_ei = './ckp/21-09-17-12:58:21_cs4x_ei/ckp_29.pth.tar' # replace the path with your trained model\n",
        "unet = UNet().to(device)\n",
        "\n",
        "# ground-truh\n",
        "x = test_dataset[0][0].type(dtype).to(device) #torch.Size([1, 1, 16, 16])\n",
        "# measurement\n",
        "y = physics.A(x)\n",
        "ATy = physics.A_dagger(y)\n",
        "x_ei = test(unet, ckp_ei, ATy)\n",
        "\n",
        "plt.subplot(1,4,1)\n",
        "plt.imshow(torch2img(y.view(1,1,4,16)), cmap='gray'), plt.title('y'), plt.axis('off')\n",
        "plt.subplot(1,4,2)\n",
        "plt.imshow(torch2img(ATy), cmap='gray'), plt.title(r'$A^+y$'+'\\n{:.2f}'.format(cal_psnr(ATy, x))), plt.axis('off')\n",
        "plt.subplot(1,4,3)\n",
        "plt.imshow(torch2img(x_ei), cmap='gray'), plt.title('EI\\n{:.2f}'.format(cal_psnr(x_ei, x))), plt.axis('off')\n",
        "plt.subplot(1,4,4)\n",
        "plt.imshow(torch2img(x), cmap='gray'), plt.title('x (GT)'), plt.axis('off')\n",
        "plt.show()"
      ],
      "execution_count": null,
      "outputs": []
    }
  ]
}