demo_train.ipynb

{
  "cells": [
    {
      "cell_type": "code",
      "source": [
        "#@title Download NTIRE202 validation dataset (toy example)\n",
        "!wget --load-cookies /tmp/cookies.txt \"https://docs.google.com/uc?export=download&confirm=$(wget --quiet --save-cookies /tmp/cookies.txt --keep-session-cookies --no-check-certificate 'https://docs.google.com/uc?export=download&id=1IK2emmr39ribjE7zeIGW221u8na8_ar2' -O- | sed -rn 's/.*confirm=([0-9A-Za-z_]+).*/\\1\\n/p')&id=1IK2emmr39ribjE7zeIGW221u8na8_ar2\" -O NTIRE2020_Validation_Spectral.zip && rm -rf /tmp/cookies.txt\n",
        "!mkdir /content/data\n",
        "!unzip NTIRE2020_Validation_Spectral -d /content/data"
      ],
      "metadata": {
        "id": "z5uArM5M7twj",
        "outputId": "9fac5290-8d8d-4c51-c849-16fe194ad853",
        "colab": {
          "base_uri": "https://localhost:8080/"
        }
      },
      "execution_count": 1,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "--2023-01-12 15:44:59--  https://docs.google.com/uc?export=download&confirm=t&id=1IK2emmr39ribjE7zeIGW221u8na8_ar2\n",
            "Resolving docs.google.com (docs.google.com)... 74.125.24.139, 74.125.24.101, 74.125.24.102, ...\n",
            "Connecting to docs.google.com (docs.google.com)|74.125.24.139|:443... connected.\n",
            "HTTP request sent, awaiting response... 303 See Other\n",
            "Location: https://doc-08-88-docs.googleusercontent.com/docs/securesc/ha0ro937gcuc7l7deffksulhg5h7mbp1/79fc7gf148b6b32arb4amsfspchjnd0m/1673538225000/14395628383364312969/*/1IK2emmr39ribjE7zeIGW221u8na8_ar2?e=download&uuid=64347988-11fa-41c3-9191-ca4d6b4e9d78 [following]\n",
            "Warning: wildcards not supported in HTTP.\n",
            "--2023-01-12 15:44:59--  https://doc-08-88-docs.googleusercontent.com/docs/securesc/ha0ro937gcuc7l7deffksulhg5h7mbp1/79fc7gf148b6b32arb4amsfspchjnd0m/1673538225000/14395628383364312969/*/1IK2emmr39ribjE7zeIGW221u8na8_ar2?e=download&uuid=64347988-11fa-41c3-9191-ca4d6b4e9d78\n",
            "Resolving doc-08-88-docs.googleusercontent.com (doc-08-88-docs.googleusercontent.com)... 74.125.24.132, 2404:6800:4003:c03::84\n",
            "Connecting to doc-08-88-docs.googleusercontent.com (doc-08-88-docs.googleusercontent.com)|74.125.24.132|:443... connected.\n",
            "HTTP request sent, awaiting response... 200 OK\n",
            "Length: 572792086 (546M) [application/zip]\n",
            "Saving to: ‘NTIRE2020_Validation_Spectral.zip’\n",
            "\n",
            "NTIRE2020_Validatio 100%[===================>] 546.26M  53.5MB/s    in 11s     \n",
            "\n",
            "2023-01-12 15:45:11 (49.1 MB/s) - ‘NTIRE2020_Validation_Spectral.zip’ saved [572792086/572792086]\n",
            "\n",
            "Archive:  NTIRE2020_Validation_Spectral.zip\n",
            "  inflating: /content/data/ARAD_HS_0451.mat  \n",
            "  inflating: /content/data/ARAD_HS_0453.mat  \n",
            "  inflating: /content/data/ARAD_HS_0455.mat  \n",
            "  inflating: /content/data/ARAD_HS_0456.mat  \n",
            "  inflating: /content/data/ARAD_HS_0457.mat  \n",
            "  inflating: /content/data/ARAD_HS_0459.mat  \n",
            "  inflating: /content/data/ARAD_HS_0462.mat  \n",
            "  inflating: /content/data/ARAD_HS_0463.mat  \n",
            "  inflating: /content/data/ARAD_HS_0464.mat  \n",
            "  inflating: /content/data/ARAD_HS_0465.mat  \n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "#@title Clone JR2net Github Repository\n",
        "!git clone https://github.com/bemc22/JR2net\n",
        "%cd JR2net"
      ],
      "metadata": {
        "cellView": "form",
        "id": "EnB0iCEM39_u",
        "outputId": "1bdc059a-898d-41b7-b3f1-ceb7da3345c4",
        "colab": {
          "base_uri": "https://localhost:8080/"
        }
      },
      "execution_count": 2,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Cloning into 'JR2net'...\n",
            "remote: Enumerating objects: 95, done.\u001b[K\n",
            "remote: Counting objects: 100% (95/95), done.\u001b[K\n",
            "remote: Compressing objects: 100% (78/78), done.\u001b[K\n",
            "remote: Total 95 (delta 44), reused 35 (delta 12), pack-reused 0\u001b[K\n",
            "Unpacking objects: 100% (95/95), done.\n",
            "/content/JR2net\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 3,
      "metadata": {
        "id": "sogRzAhN3WxW"
      },
      "outputs": [],
      "source": [
        "import os\n",
        "import tensorflow as tf\n",
        "import numpy as np\n",
        "import scipy.io as sio\n",
        "import matplotlib.pyplot as plt\n",
        "\n",
        "from dataset import get_csi_pipeline, get_val_csi\n",
        "from jr2net.utils import dd_cassi , dd_icassi, coded2DTO3D"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 4,
      "metadata": {
        "id": "4rkWwf8X3WxZ"
      },
      "outputs": [],
      "source": [
        "RGB = [27, 17 , 4] # RGB channels\n",
        "BATCH_SIZE = 3     # Amount of images samples\n",
        "PATCHES = 9        # The final batch size is BATCH_SIZE * PATCHES\n",
        "split = 0.9\n",
        "INPUT_SHAPE = (104, 104, 31)\n",
        "DATASE_SIZE = (482, 512, 31)\n",
        "\n",
        "FOLDER_PATH = \"/content\"\n",
        "\n",
        "data_path = os.path.join(FOLDER_PATH, 'data')\n",
        "test_path = data_path\n",
        "\n",
        "DATA__PARAMS = dict(\n",
        "    input_size = INPUT_SHAPE,\n",
        "    batch_size = BATCH_SIZE,\n",
        "    origin_size = DATASE_SIZE,\n",
        ")\n",
        "\n",
        "\n",
        "train_ds = get_csi_pipeline(data_path=data_path, buffer_size=10, **DATA__PARAMS, factor=PATCHES)\n",
        "# val_ds = get_csi_pipeline(data_path=test_path, buffer_size=None, **DATA__PARAMS)\n",
        "val_ds = get_val_csi(data_path)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 5,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "GtEko1Pi3Wxa",
        "outputId": "7c8f391b-ee91-4572-af58-bca59fff9767"
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Model: \"JR2net\"\n",
            "__________________________________________________________________________________________________\n",
            " Layer (type)                   Output Shape         Param #     Connected to                     \n",
            "==================================================================================================\n",
            " spectral_image (InputLayer)    [(None, None, None,  0           []                               \n",
            "                                 31)]                                                             \n",
            "                                                                                                  \n",
            " coded_aperture (InputLayer)    [(None, None, None,  0           []                               \n",
            "                                 31)]                                                             \n",
            "                                                                                                  \n",
            " forward_cassi (Lambda)         (None, None, None,   0           ['spectral_image[0][0]',         \n",
            "                                1)                                'coded_aperture[0][0]',         \n",
            "                                                                  'inverse_cassi[0][0]',          \n",
            "                                                                  'coded_aperture[0][0]',         \n",
            "                                                                  'Decoder[1][0]',                \n",
            "                                                                  'coded_aperture[0][0]',         \n",
            "                                                                  'Decoder[2][0]',                \n",
            "                                                                  'coded_aperture[0][0]',         \n",
            "                                                                  'Decoder[3][0]',                \n",
            "                                                                  'coded_aperture[0][0]',         \n",
            "                                                                  'Decoder[4][0]',                \n",
            "                                                                  'coded_aperture[0][0]',         \n",
            "                                                                  'Decoder[5][0]',                \n",
            "                                                                  'coded_aperture[0][0]',         \n",
            "                                                                  'Decoder[6][0]',                \n",
            "                                                                  'coded_aperture[0][0]']         \n",
            "                                                                                                  \n",
            " inverse_cassi (Lambda)         (None, None, None,   0           ['forward_cassi[0][0]',          \n",
            "                                31)                               'coded_aperture[0][0]',         \n",
            "                                                                  'forward_cassi[1][0]',          \n",
            "                                                                  'coded_aperture[0][0]',         \n",
            "                                                                  'forward_cassi[2][0]',          \n",
            "                                                                  'coded_aperture[0][0]',         \n",
            "                                                                  'forward_cassi[3][0]',          \n",
            "                                                                  'coded_aperture[0][0]',         \n",
            "                                                                  'forward_cassi[4][0]',          \n",
            "                                                                  'coded_aperture[0][0]',         \n",
            "                                                                  'forward_cassi[5][0]',          \n",
            "                                                                  'coded_aperture[0][0]',         \n",
            "                                                                  'forward_cassi[6][0]',          \n",
            "                                                                  'coded_aperture[0][0]',         \n",
            "                                                                  'forward_cassi[7][0]',          \n",
            "                                                                  'coded_aperture[0][0]']         \n",
            "                                                                                                  \n",
            " lambda (Lambda)                (None, None, None,   0           ['inverse_cassi[0][0]']          \n",
            "                                31)                                                               \n",
            "                                                                                                  \n",
            " tf.__operators__.add (TFOpLamb  (None, None, None,   0          ['inverse_cassi[0][0]',          \n",
            " da)                            31)                               'lambda[0][0]']                 \n",
            "                                                                                                  \n",
            " PriorLayer_0 (PriorLayer)      (None, None, None,   18352       ['tf.__operators__.add[0][0]']   \n",
            "                                31)                                                               \n",
            "                                                                                                  \n",
            " tf.math.subtract_1 (TFOpLambda  (None, None, None,   0          ['inverse_cassi[0][0]',          \n",
            " )                              31)                               'PriorLayer_0[0][0]']           \n",
            "                                                                                                  \n",
            " tf.math.subtract (TFOpLambda)  (None, None, None,   0           ['inverse_cassi[1][0]',          \n",
            "                                31)                               'inverse_cassi[0][0]']          \n",
            "                                                                                                  \n",
            " tf.__operators__.add_1 (TFOpLa  (None, None, None,   0          ['tf.math.subtract_1[0][0]',     \n",
            " mbda)                          31)                               'lambda[0][0]']                 \n",
            "                                                                                                  \n",
            " Encoder (EncodeLayer)          (None, None, None,   88344       ['spectral_image[0][0]',         \n",
            "                                8)                                'inverse_cassi[0][0]',          \n",
            "                                                                  'tf.math.subtract[0][0]',       \n",
            "                                                                  'tf.__operators__.add_1[0][0]', \n",
            "                                                                  'tf.math.subtract_3[0][0]',     \n",
            "                                                                  'tf.__operators__.add_4[0][0]', \n",
            "                                                                  'tf.math.subtract_6[0][0]',     \n",
            "                                                                  'tf.__operators__.add_7[0][0]', \n",
            "                                                                  'tf.math.subtract_9[0][0]',     \n",
            "                                                                  'tf.__operators__.add_10[0][0]',\n",
            "                                                                  'tf.math.subtract_12[0][0]',    \n",
            "                                                                  'tf.__operators__.add_13[0][0]',\n",
            "                                                                  'tf.math.subtract_15[0][0]',    \n",
            "                                                                  'tf.__operators__.add_16[0][0]',\n",
            "                                                                  'tf.math.subtract_18[0][0]',    \n",
            "                                                                  'tf.__operators__.add_19[0][0]']\n",
            "                                                                                                  \n",
            " GradientDescent_0 (GradientDes  (None, None, None,   2          ['Encoder[1][0]',                \n",
            " cent)                          8)                                'Encoder[2][0]',                \n",
            "                                                                  'Encoder[3][0]']                \n",
            "                                                                                                  \n",
            " Decoder (DecodeLayer)          (None, None, None,   88367       ['Encoder[0][0]',                \n",
            "                                31)                               'GradientDescent_0[0][0]',      \n",
            "                                                                  'GradientDescent_1[0][0]',      \n",
            "                                                                  'GradientDescent_2[0][0]',      \n",
            "                                                                  'GradientDescent_3[0][0]',      \n",
            "                                                                  'GradientDescent_4[0][0]',      \n",
            "                                                                  'GradientDescent_5[0][0]',      \n",
            "                                                                  'GradientDescent_6[0][0]']      \n",
            "                                                                                                  \n",
            " tf.__operators__.add_2 (TFOpLa  (None, None, None,   0          ['lambda[0][0]',                 \n",
            " mbda)                          31)                               'Decoder[1][0]']                \n",
            "                                                                                                  \n",
            " tf.math.subtract_2 (TFOpLambda  (None, None, None,   0          ['tf.__operators__.add_2[0][0]', \n",
            " )                              31)                               'PriorLayer_0[0][0]']           \n",
            "                                                                                                  \n",
            " tf.__operators__.add_3 (TFOpLa  (None, None, None,   0          ['Decoder[1][0]',                \n",
            " mbda)                          31)                               'tf.math.subtract_2[0][0]']     \n",
            "                                                                                                  \n",
            " PriorLayer_1 (PriorLayer)      (None, None, None,   18352       ['tf.__operators__.add_3[0][0]'] \n",
            "                                31)                                                               \n",
            "                                                                                                  \n",
            " tf.math.subtract_4 (TFOpLambda  (None, None, None,   0          ['Decoder[1][0]',                \n",
            " )                              31)                               'PriorLayer_1[0][0]']           \n",
            "                                                                                                  \n",
            " tf.math.subtract_3 (TFOpLambda  (None, None, None,   0          ['inverse_cassi[2][0]',          \n",
            " )                              31)                               'inverse_cassi[0][0]']          \n",
            "                                                                                                  \n",
            " tf.__operators__.add_4 (TFOpLa  (None, None, None,   0          ['tf.math.subtract_4[0][0]',     \n",
            " mbda)                          31)                               'tf.math.subtract_2[0][0]']     \n",
            "                                                                                                  \n",
            " GradientDescent_1 (GradientDes  (None, None, None,   2          ['GradientDescent_0[0][0]',      \n",
            " cent)                          8)                                'Encoder[4][0]',                \n",
            "                                                                  'Encoder[5][0]']                \n",
            "                                                                                                  \n",
            " tf.__operators__.add_5 (TFOpLa  (None, None, None,   0          ['tf.math.subtract_2[0][0]',     \n",
            " mbda)                          31)                               'Decoder[2][0]']                \n",
            "                                                                                                  \n",
            " tf.math.subtract_5 (TFOpLambda  (None, None, None,   0          ['tf.__operators__.add_5[0][0]', \n",
            " )                              31)                               'PriorLayer_1[0][0]']           \n",
            "                                                                                                  \n",
            " tf.__operators__.add_6 (TFOpLa  (None, None, None,   0          ['Decoder[2][0]',                \n",
            " mbda)                          31)                               'tf.math.subtract_5[0][0]']     \n",
            "                                                                                                  \n",
            " PriorLayer_2 (PriorLayer)      (None, None, None,   18352       ['tf.__operators__.add_6[0][0]'] \n",
            "                                31)                                                               \n",
            "                                                                                                  \n",
            " tf.math.subtract_7 (TFOpLambda  (None, None, None,   0          ['Decoder[2][0]',                \n",
            " )                              31)                               'PriorLayer_2[0][0]']           \n",
            "                                                                                                  \n",
            " tf.math.subtract_6 (TFOpLambda  (None, None, None,   0          ['inverse_cassi[3][0]',          \n",
            " )                              31)                               'inverse_cassi[0][0]']          \n",
            "                                                                                                  \n",
            " tf.__operators__.add_7 (TFOpLa  (None, None, None,   0          ['tf.math.subtract_7[0][0]',     \n",
            " mbda)                          31)                               'tf.math.subtract_5[0][0]']     \n",
            "                                                                                                  \n",
            " GradientDescent_2 (GradientDes  (None, None, None,   2          ['GradientDescent_1[0][0]',      \n",
            " cent)                          8)                                'Encoder[6][0]',                \n",
            "                                                                  'Encoder[7][0]']                \n",
            "                                                                                                  \n",
            " tf.__operators__.add_8 (TFOpLa  (None, None, None,   0          ['tf.math.subtract_5[0][0]',     \n",
            " mbda)                          31)                               'Decoder[3][0]']                \n",
            "                                                                                                  \n",
            " tf.math.subtract_8 (TFOpLambda  (None, None, None,   0          ['tf.__operators__.add_8[0][0]', \n",
            " )                              31)                               'PriorLayer_2[0][0]']           \n",
            "                                                                                                  \n",
            " tf.__operators__.add_9 (TFOpLa  (None, None, None,   0          ['Decoder[3][0]',                \n",
            " mbda)                          31)                               'tf.math.subtract_8[0][0]']     \n",
            "                                                                                                  \n",
            " PriorLayer_3 (PriorLayer)      (None, None, None,   18352       ['tf.__operators__.add_9[0][0]'] \n",
            "                                31)                                                               \n",
            "                                                                                                  \n",
            " tf.math.subtract_10 (TFOpLambd  (None, None, None,   0          ['Decoder[3][0]',                \n",
            " a)                             31)                               'PriorLayer_3[0][0]']           \n",
            "                                                                                                  \n",
            " tf.math.subtract_9 (TFOpLambda  (None, None, None,   0          ['inverse_cassi[4][0]',          \n",
            " )                              31)                               'inverse_cassi[0][0]']          \n",
            "                                                                                                  \n",
            " tf.__operators__.add_10 (TFOpL  (None, None, None,   0          ['tf.math.subtract_10[0][0]',    \n",
            " ambda)                         31)                               'tf.math.subtract_8[0][0]']     \n",
            "                                                                                                  \n",
            " GradientDescent_3 (GradientDes  (None, None, None,   2          ['GradientDescent_2[0][0]',      \n",
            " cent)                          8)                                'Encoder[8][0]',                \n",
            "                                                                  'Encoder[9][0]']                \n",
            "                                                                                                  \n",
            " tf.__operators__.add_11 (TFOpL  (None, None, None,   0          ['tf.math.subtract_8[0][0]',     \n",
            " ambda)                         31)                               'Decoder[4][0]']                \n",
            "                                                                                                  \n",
            " tf.math.subtract_11 (TFOpLambd  (None, None, None,   0          ['tf.__operators__.add_11[0][0]',\n",
            " a)                             31)                               'PriorLayer_3[0][0]']           \n",
            "                                                                                                  \n",
            " tf.__operators__.add_12 (TFOpL  (None, None, None,   0          ['Decoder[4][0]',                \n",
            " ambda)                         31)                               'tf.math.subtract_11[0][0]']    \n",
            "                                                                                                  \n",
            " PriorLayer_4 (PriorLayer)      (None, None, None,   18352       ['tf.__operators__.add_12[0][0]']\n",
            "                                31)                                                               \n",
            "                                                                                                  \n",
            " tf.math.subtract_13 (TFOpLambd  (None, None, None,   0          ['Decoder[4][0]',                \n",
            " a)                             31)                               'PriorLayer_4[0][0]']           \n",
            "                                                                                                  \n",
            " tf.math.subtract_12 (TFOpLambd  (None, None, None,   0          ['inverse_cassi[5][0]',          \n",
            " a)                             31)                               'inverse_cassi[0][0]']          \n",
            "                                                                                                  \n",
            " tf.__operators__.add_13 (TFOpL  (None, None, None,   0          ['tf.math.subtract_13[0][0]',    \n",
            " ambda)                         31)                               'tf.math.subtract_11[0][0]']    \n",
            "                                                                                                  \n",
            " GradientDescent_4 (GradientDes  (None, None, None,   2          ['GradientDescent_3[0][0]',      \n",
            " cent)                          8)                                'Encoder[10][0]',               \n",
            "                                                                  'Encoder[11][0]']               \n",
            "                                                                                                  \n",
            " tf.__operators__.add_14 (TFOpL  (None, None, None,   0          ['tf.math.subtract_11[0][0]',    \n",
            " ambda)                         31)                               'Decoder[5][0]']                \n",
            "                                                                                                  \n",
            " tf.math.subtract_14 (TFOpLambd  (None, None, None,   0          ['tf.__operators__.add_14[0][0]',\n",
            " a)                             31)                               'PriorLayer_4[0][0]']           \n",
            "                                                                                                  \n",
            " tf.__operators__.add_15 (TFOpL  (None, None, None,   0          ['Decoder[5][0]',                \n",
            " ambda)                         31)                               'tf.math.subtract_14[0][0]']    \n",
            "                                                                                                  \n",
            " PriorLayer_5 (PriorLayer)      (None, None, None,   18352       ['tf.__operators__.add_15[0][0]']\n",
            "                                31)                                                               \n",
            "                                                                                                  \n",
            " tf.math.subtract_16 (TFOpLambd  (None, None, None,   0          ['Decoder[5][0]',                \n",
            " a)                             31)                               'PriorLayer_5[0][0]']           \n",
            "                                                                                                  \n",
            " tf.math.subtract_15 (TFOpLambd  (None, None, None,   0          ['inverse_cassi[6][0]',          \n",
            " a)                             31)                               'inverse_cassi[0][0]']          \n",
            "                                                                                                  \n",
            " tf.__operators__.add_16 (TFOpL  (None, None, None,   0          ['tf.math.subtract_16[0][0]',    \n",
            " ambda)                         31)                               'tf.math.subtract_14[0][0]']    \n",
            "                                                                                                  \n",
            " GradientDescent_5 (GradientDes  (None, None, None,   2          ['GradientDescent_4[0][0]',      \n",
            " cent)                          8)                                'Encoder[12][0]',               \n",
            "                                                                  'Encoder[13][0]']               \n",
            "                                                                                                  \n",
            " tf.__operators__.add_17 (TFOpL  (None, None, None,   0          ['tf.math.subtract_14[0][0]',    \n",
            " ambda)                         31)                               'Decoder[6][0]']                \n",
            "                                                                                                  \n",
            " tf.math.subtract_17 (TFOpLambd  (None, None, None,   0          ['tf.__operators__.add_17[0][0]',\n",
            " a)                             31)                               'PriorLayer_5[0][0]']           \n",
            "                                                                                                  \n",
            " tf.__operators__.add_18 (TFOpL  (None, None, None,   0          ['Decoder[6][0]',                \n",
            " ambda)                         31)                               'tf.math.subtract_17[0][0]']    \n",
            "                                                                                                  \n",
            " PriorLayer_6 (PriorLayer)      (None, None, None,   18352       ['tf.__operators__.add_18[0][0]']\n",
            "                                31)                                                               \n",
            "                                                                                                  \n",
            " tf.math.subtract_19 (TFOpLambd  (None, None, None,   0          ['Decoder[6][0]',                \n",
            " a)                             31)                               'PriorLayer_6[0][0]']           \n",
            "                                                                                                  \n",
            " tf.math.subtract_18 (TFOpLambd  (None, None, None,   0          ['inverse_cassi[7][0]',          \n",
            " a)                             31)                               'inverse_cassi[0][0]']          \n",
            "                                                                                                  \n",
            " tf.__operators__.add_19 (TFOpL  (None, None, None,   0          ['tf.math.subtract_19[0][0]',    \n",
            " ambda)                         31)                               'tf.math.subtract_17[0][0]']    \n",
            "                                                                                                  \n",
            " GradientDescent_6 (GradientDes  (None, None, None,   2          ['GradientDescent_5[0][0]',      \n",
            " cent)                          8)                                'Encoder[14][0]',               \n",
            "                                                                  'Encoder[15][0]']               \n",
            "                                                                                                  \n",
            " recons (Lambda)                (None, None, None,   0           ['Decoder[7][0]']                \n",
            "                                31)                                                               \n",
            "                                                                                                  \n",
            " repre (Lambda)                 (None, None, None,   0           ['Decoder[0][0]']                \n",
            "                                31)                                                               \n",
            "                                                                                                  \n",
            "==================================================================================================\n",
            "Total params: 305,189\n",
            "Trainable params: 305,189\n",
            "Non-trainable params: 0\n",
            "__________________________________________________________________________________________________\n"
          ]
        }
      ],
      "source": [
        "from gc import callbacks\n",
        "from jr2net.models import JR2net\n",
        "from jr2net.metrics import prior_loss, psnr\n",
        "from jr2net.utils import ClearMemory\n",
        "from tensorflow.keras.callbacks import ModelCheckpoint, Callback, ReduceLROnPlateau\n",
        "from tensorflow.keras import backend as k\n",
        "\n",
        "STAGES = 7\n",
        "FACTORS = [1, 1, 1/2, 1/2, 1/4, 1/8]\n",
        "EPOCHS = 10\n",
        "\n",
        "MODEL_PARAMS = {\n",
        "    'loss': ['mse', 'mse'],\n",
        "    'optimizer': tf.keras.optimizers.Adam(learning_rate=1e-4, amsgrad=False),\n",
        "    'metrics': [psnr]\n",
        "}\n",
        "\n",
        "\n",
        "UNROLLED_PARAMS = {\n",
        "    'input_size' : (None, None, INPUT_SHAPE[-1]),\n",
        "    'num_stages': STAGES,\n",
        "    'factors': FACTORS,\n",
        "    'training': True,\n",
        "}\n",
        "\n",
        "unrolled_weights = \"unrolled_weights.h5\"\n",
        "callbacks = [ModelCheckpoint(unrolled_weights,\n",
        "                             monitor='recons_psnr',\n",
        "                             save_best_only=True,\n",
        "                             save_weights_only=True,\n",
        "                             mode='max'),\n",
        "             ClearMemory(),\n",
        "             ReduceLROnPlateau(monitor=\"recons_psnr\",\n",
        "                               factor=0.9,\n",
        "                               patience=100,\n",
        "                               verbose=0,\n",
        "                               mode=\"max\")]\n",
        "\n",
        "main_model = JR2net(**UNROLLED_PARAMS)\n",
        "model = main_model.unrolled\n",
        "model.compile(**MODEL_PARAMS)\n",
        "model.summary()\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 6,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "2cacx4013Wxc",
        "outputId": "d368d155-9a40-44d3-d797-f86c638b3367"
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Epoch 1/10\n",
            "3/3 [==============================] - 23s 2s/step - loss: 0.1536 - recons_loss: 0.0768 - repre_loss: 0.0768 - recons_psnr: 13.8206 - repre_psnr: 13.8227\n",
            "Epoch 2/10\n",
            "3/3 [==============================] - 6s 2s/step - loss: 0.1495 - recons_loss: 0.0750 - repre_loss: 0.0745 - recons_psnr: 13.9074 - repre_psnr: 13.9362\n",
            "Epoch 3/10\n",
            "3/3 [==============================] - 6s 2s/step - loss: 0.1418 - recons_loss: 0.0715 - repre_loss: 0.0703 - recons_psnr: 13.7764 - repre_psnr: 13.8544\n",
            "Epoch 4/10\n",
            "3/3 [==============================] - 6s 2s/step - loss: 0.1223 - recons_loss: 0.0623 - repre_loss: 0.0601 - recons_psnr: 13.9160 - repre_psnr: 14.0827\n",
            "Epoch 5/10\n",
            "3/3 [==============================] - 6s 2s/step - loss: 0.1302 - recons_loss: 0.0673 - repre_loss: 0.0628 - recons_psnr: 14.6824 - repre_psnr: 14.9917\n",
            "Epoch 6/10\n",
            "3/3 [==============================] - 6s 2s/step - loss: 0.1117 - recons_loss: 0.0592 - repre_loss: 0.0525 - recons_psnr: 14.5272 - repre_psnr: 15.0490\n",
            "Epoch 7/10\n",
            "3/3 [==============================] - 6s 2s/step - loss: 0.0914 - recons_loss: 0.0492 - repre_loss: 0.0422 - recons_psnr: 15.3846 - repre_psnr: 15.9715\n",
            "Epoch 8/10\n",
            "3/3 [==============================] - 6s 2s/step - loss: 0.0814 - recons_loss: 0.0397 - repre_loss: 0.0417 - recons_psnr: 16.0272 - repre_psnr: 15.4308\n",
            "Epoch 9/10\n",
            "3/3 [==============================] - 6s 2s/step - loss: 0.0905 - recons_loss: 0.0415 - repre_loss: 0.0490 - recons_psnr: 16.2276 - repre_psnr: 15.1109\n",
            "Epoch 10/10\n",
            "3/3 [==============================] - 6s 2s/step - loss: 0.0793 - recons_loss: 0.0400 - repre_loss: 0.0392 - recons_psnr: 16.4491 - repre_psnr: 16.2367\n"
          ]
        },
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "<keras.callbacks.History at 0x7fd37a707c10>"
            ]
          },
          "metadata": {},
          "execution_count": 6
        }
      ],
      "source": [
        "model.fit(train_ds, epochs=EPOCHS, validation_data=None, callbacks=None)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 6,
      "metadata": {
        "id": "cEmO62O23Wxc"
      },
      "outputs": [],
      "source": []
    }
  ],
  "metadata": {
    "kernelspec": {
      "display_name": "Python 3.9.6 ('deep-learning')",
      "language": "python",
      "name": "python3"
    },
    "language_info": {
      "codemirror_mode": {
        "name": "ipython",
        "version": 3
      },
      "file_extension": ".py",
      "mimetype": "text/x-python",
      "name": "python",
      "nbconvert_exporter": "python",
      "pygments_lexer": "ipython3",
      "version": "3.9.6"
    },
    "orig_nbformat": 4,
    "vscode": {
      "interpreter": {
        "hash": "e19c5b7ce74c29a4a0f33606fce46481bd4d1c925c075dd8e930b49a84ff9795"
      }
    },
    "colab": {
      "provenance": []
    },
    "accelerator": "GPU",
    "gpuClass": "standard"
  },
  "nbformat": 4,
  "nbformat_minor": 0
}