test_unet.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "9b510c10-e4eb-471d-b9a9-d79efe1c8e16",
   "metadata": {},
   "source": [
    "# Test U-Net"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7b5ed535-28c9-4ef7-bff4-ed872a4dbebf",
   "metadata": {},
   "source": [
    "## U-Net architecture"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "0d8e7aac-d5e3-4ace-8078-3cbebcde32b0",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "2023-04-13 10:34:13.647641: I tensorflow/core/platform/cpu_feature_guard.cc:193] This TensorFlow binary is optimized with oneAPI Deep Neural Network Library (oneDNN) to use the following CPU instructions in performance-critical operations:  AVX2 AVX512F AVX512_VNNI FMA\n",
      "To enable them in other operations, rebuild TensorFlow with the appropriate compiler flags.\n",
      "2023-04-13 10:34:13.743638: I tensorflow/core/util/util.cc:169] oneDNN custom operations are on. You may see slightly different numerical results due to floating-point round-off errors from different computation orders. To turn them off, set the environment variable `TF_ENABLE_ONEDNN_OPTS=0`.\n",
      "2023-04-13 10:34:13.767841: E tensorflow/stream_executor/cuda/cuda_blas.cc:2981] Unable to register cuBLAS factory: Attempting to register factory for plugin cuBLAS when one has already been registered\n",
      "/etinfo/users2/qlanglois/.local/lib/python3.9/site-packages/scipy/__init__.py:146: UserWarning: A NumPy version >=1.16.5 and <1.23.0 is required for this version of SciPy (detected version 1.23.4\n",
      "  warnings.warn(f\"A NumPy version >={np_minversion} and <{np_maxversion}\"\n",
      "/etinfo/users2/qlanglois/.local/lib/python3.9/site-packages/tensorflow_addons/utils/tfa_eol_msg.py:23: UserWarning: \n",
      "\n",
      "TensorFlow Addons (TFA) has ended development and introduction of new features.\n",
      "TFA has entered a minimal maintenance and release mode until a planned end of life in May 2024.\n",
      "Please modify downstream libraries to take dependencies from other repositories in our TensorFlow community (e.g. Keras, Keras-CV, and Keras-NLP). \n",
      "\n",
      "For more information see: https://github.com/tensorflow/addons/issues/2807 \n",
      "\n",
      "  warnings.warn(\n"
     ]
    }
   ],
   "source": [
    "import importlib\n",
    "import numpy as np\n",
    "import tensorflow as tf\n",
    "\n",
    "from custom_architectures import unet_arch, totalsegmentator_arch\n",
    "\n",
    "def show_memory(msg = ''):\n",
    "    print('{}{}'.format(msg if not msg else msg + '\\t: ', {\n",
    "        k : '{:.3f}'.format(v / 1024 ** 3) for k, v in tf.config.experimental.get_memory_info('GPU:0').items()\n",
    "    }))\n",
    "    tf.config.experimental.reset_memory_stats('GPU:0')\n",
    "    \n",
    "gpus = tf.config.list_physical_devices('GPU')\n",
    "tf.config.set_visible_devices([gpus[0]], 'GPU')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "daea3419-ee8e-4c4c-8264-2278aef71a77",
   "metadata": {
    "jp-MarkdownHeadingCollapsed": true,
    "tags": []
   },
   "source": [
    "## Test U-Net 2D"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "99a16581-d03c-49b6-93b9-16fc990f1d8b",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "2023-03-23 15:35:39.337311: I tensorflow/core/platform/cpu_feature_guard.cc:193] This TensorFlow binary is optimized with oneAPI Deep Neural Network Library (oneDNN) to use the following CPU instructions in performance-critical operations:  AVX2 AVX512F AVX512_VNNI FMA\n",
      "To enable them in other operations, rebuild TensorFlow with the appropriate compiler flags.\n",
      "2023-03-23 15:35:39.707179: I tensorflow/core/common_runtime/gpu/gpu_device.cc:1616] Created device /job:localhost/replica:0/task:0/device:GPU:0 with 14783 MB memory:  -> device: 0, name: Quadro RTX 5000, pci bus id: 0000:17:00.0, compute capability: 7.5\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "8611425\n",
      "{'current': '0.032', 'peak': '0.035'}\n"
     ]
    }
   ],
   "source": [
    "importlib.reload(unet_arch)\n",
    "\n",
    "model = unet_arch.UNet(\n",
    "    input_shape = (512, 512, 1), output_dim  = 1,\n",
    "    drop_rate = lambda i: 0. if i == 4 else 0.25\n",
    ")\n",
    "print(model.count_params())\n",
    "model.compile(loss = 'binary_crossentropy', optimizer = 'adam')\n",
    "\n",
    "show_memory()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "4c6293fb-f2e6-4494-b184-bb47fd88fe3f",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Before\t: {'current': '0.032', 'peak': '0.035'}\n",
      "Epoch 1/5\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "2023-03-23 15:35:41.686235: I tensorflow/stream_executor/cuda/cuda_dnn.cc:384] Loaded cuDNN version 8204\n",
      "2023-03-23 15:35:49.280779: W tensorflow/core/common_runtime/bfc_allocator.cc:290] Allocator (GPU_0_bfc) ran out of memory trying to allocate 5.08GiB with freed_by_count=0. The caller indicates that this is not a failure, but this may mean that there could be performance gains if more memory were available.\n",
      "2023-03-23 15:35:49.280808: W tensorflow/core/common_runtime/bfc_allocator.cc:290] Allocator (GPU_0_bfc) ran out of memory trying to allocate 5.08GiB with freed_by_count=0. The caller indicates that this is not a failure, but this may mean that there could be performance gains if more memory were available.\n",
      "2023-03-23 15:35:49.683526: W tensorflow/core/common_runtime/bfc_allocator.cc:290] Allocator (GPU_0_bfc) ran out of memory trying to allocate 1.52GiB with freed_by_count=0. The caller indicates that this is not a failure, but this may mean that there could be performance gains if more memory were available.\n",
      "2023-03-23 15:35:49.683549: W tensorflow/core/kernels/gpu_utils.cc:50] Failed to allocate memory for convolution redzone checking; skipping this check. This is benign and only means that we won't check cudnn for out-of-bounds reads and writes. This message will only be printed once.\n",
      "2023-03-23 15:35:49.683936: W tensorflow/core/common_runtime/bfc_allocator.cc:290] Allocator (GPU_0_bfc) ran out of memory trying to allocate 1.52GiB with freed_by_count=0. The caller indicates that this is not a failure, but this may mean that there could be performance gains if more memory were available.\n",
      "2023-03-23 15:35:49.683952: W tensorflow/core/common_runtime/bfc_allocator.cc:290] Allocator (GPU_0_bfc) ran out of memory trying to allocate 1.52GiB with freed_by_count=0. The caller indicates that this is not a failure, but this may mean that there could be performance gains if more memory were available.\n",
      "2023-03-23 15:35:49.876063: W tensorflow/core/common_runtime/bfc_allocator.cc:290] Allocator (GPU_0_bfc) ran out of memory trying to allocate 6.77GiB with freed_by_count=0. The caller indicates that this is not a failure, but this may mean that there could be performance gains if more memory were available.\n",
      "2023-03-23 15:35:49.876085: W tensorflow/core/common_runtime/bfc_allocator.cc:290] Allocator (GPU_0_bfc) ran out of memory trying to allocate 6.77GiB with freed_by_count=0. The caller indicates that this is not a failure, but this may mean that there could be performance gains if more memory were available.\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "2/2 [==============================] - 16s 899ms/step - loss: 0.6594\n",
      "Epoch 2/5\n",
      "2/2 [==============================] - 2s 900ms/step - loss: 0.1900\n",
      "Epoch 3/5\n",
      "2/2 [==============================] - 2s 898ms/step - loss: 0.0013\n",
      "Epoch 4/5\n",
      "2/2 [==============================] - 2s 897ms/step - loss: 7.1392e-05\n",
      "Epoch 5/5\n",
      "2/2 [==============================] - 2s 896ms/step - loss: 0.0000e+00\n",
      "After fit\t: {'current': '0.219', 'peak': '13.966'}\n"
     ]
    }
   ],
   "source": [
    "batch_size = 24\n",
    "\n",
    "show_memory('Before')\n",
    "\n",
    "inp = tf.random.uniform((batch_size * 2, ) + tuple(model.input_shape[1:]), 0., 1.)\n",
    "out = tf.ones_like(inp)\n",
    "\n",
    "\"\"\"show_memory('Batch init')\n",
    "\n",
    "model(inp[:batch_size], training = True)\n",
    "\n",
    "show_memory('Simple call')\n",
    "\n",
    "with tf.GradientTape() as tape:\n",
    "    pred = model(inp[:batch_size], training = True)\n",
    "    l = model.compiled_loss(out[:batch_size], pred)\n",
    "\n",
    "show_memory('With gradient')\n",
    "\n",
    "del tape, l, pred\n",
    "\n",
    "show_memory('After deleting tape')\n",
    "\"\"\"\n",
    "_ = model.fit(inp, out, epochs = 5, batch_size = batch_size)\n",
    "\n",
    "show_memory('After fit')\n",
    "tf.config.experimental.reset_memory_stats('GPU:0')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "a73975b3-e9a4-47cc-be2c-ccdad49eb076",
   "metadata": {
    "collapsed": true,
    "jupyter": {
     "outputs_hidden": true
    },
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Model: \"model\"\n",
      "__________________________________________________________________________________________________\n",
      " Layer (type)                   Output Shape         Param #     Connected to                     \n",
      "==================================================================================================\n",
      " input_image (InputLayer)       [(None, 512, 512, 1  0           []                               \n",
      "                                )]                                                                \n",
      "                                                                                                  \n",
      " down_conv1 (Conv2D)            (None, 512, 512, 32  320         ['input_image[0][0]']            \n",
      "                                )                                                                 \n",
      "                                                                                                  \n",
      " activation (Activation)        (None, 512, 512, 32  0           ['down_conv1[0][0]']             \n",
      "                                )                                                                 \n",
      "                                                                                                  \n",
      " max_pooling2d (MaxPooling2D)   (None, 256, 256, 32  0           ['activation[0][0]']             \n",
      "                                )                                                                 \n",
      "                                                                                                  \n",
      " dropout (Dropout)              (None, 256, 256, 32  0           ['max_pooling2d[0][0]']          \n",
      "                                )                                                                 \n",
      "                                                                                                  \n",
      " down_conv2-1 (Conv2D)          (None, 256, 256, 64  18496       ['dropout[0][0]']                \n",
      "                                )                                                                 \n",
      "                                                                                                  \n",
      " activation_1 (Activation)      (None, 256, 256, 64  0           ['down_conv2-1[0][0]']           \n",
      "                                )                                                                 \n",
      "                                                                                                  \n",
      " down_conv2-2 (Conv2D)          (None, 256, 256, 64  36928       ['activation_1[0][0]']           \n",
      "                                )                                                                 \n",
      "                                                                                                  \n",
      " activation_2 (Activation)      (None, 256, 256, 64  0           ['down_conv2-2[0][0]']           \n",
      "                                )                                                                 \n",
      "                                                                                                  \n",
      " max_pooling2d_1 (MaxPooling2D)  (None, 128, 128, 64  0          ['activation_2[0][0]']           \n",
      "                                )                                                                 \n",
      "                                                                                                  \n",
      " dropout_1 (Dropout)            (None, 128, 128, 64  0           ['max_pooling2d_1[0][0]']        \n",
      "                                )                                                                 \n",
      "                                                                                                  \n",
      " down_conv3-1 (Conv2D)          (None, 128, 128, 12  73856       ['dropout_1[0][0]']              \n",
      "                                8)                                                                \n",
      "                                                                                                  \n",
      " activation_3 (Activation)      (None, 128, 128, 12  0           ['down_conv3-1[0][0]']           \n",
      "                                8)                                                                \n",
      "                                                                                                  \n",
      " down_conv3-2 (Conv2D)          (None, 128, 128, 12  147584      ['activation_3[0][0]']           \n",
      "                                8)                                                                \n",
      "                                                                                                  \n",
      " activation_4 (Activation)      (None, 128, 128, 12  0           ['down_conv3-2[0][0]']           \n",
      "                                8)                                                                \n",
      "                                                                                                  \n",
      " max_pooling2d_2 (MaxPooling2D)  (None, 64, 64, 128)  0          ['activation_4[0][0]']           \n",
      "                                                                                                  \n",
      " dropout_2 (Dropout)            (None, 64, 64, 128)  0           ['max_pooling2d_2[0][0]']        \n",
      "                                                                                                  \n",
      " down_conv4-1 (Conv2D)          (None, 64, 64, 256)  295168      ['dropout_2[0][0]']              \n",
      "                                                                                                  \n",
      " activation_5 (Activation)      (None, 64, 64, 256)  0           ['down_conv4-1[0][0]']           \n",
      "                                                                                                  \n",
      " down_conv4-2 (Conv2D)          (None, 64, 64, 256)  590080      ['activation_5[0][0]']           \n",
      "                                                                                                  \n",
      " activation_6 (Activation)      (None, 64, 64, 256)  0           ['down_conv4-2[0][0]']           \n",
      "                                                                                                  \n",
      " max_pooling2d_3 (MaxPooling2D)  (None, 32, 32, 256)  0          ['activation_6[0][0]']           \n",
      "                                                                                                  \n",
      " dropout_3 (Dropout)            (None, 32, 32, 256)  0           ['max_pooling2d_3[0][0]']        \n",
      "                                                                                                  \n",
      " middle_conv1-1 (Conv2D)        (None, 32, 32, 512)  1180160     ['dropout_3[0][0]']              \n",
      "                                                                                                  \n",
      " activation_7 (Activation)      (None, 32, 32, 512)  0           ['middle_conv1-1[0][0]']         \n",
      "                                                                                                  \n",
      " dropout_4 (Dropout)            (None, 32, 32, 512)  0           ['activation_7[0][0]']           \n",
      "                                                                                                  \n",
      " middle_conv1-2 (Conv2D)        (None, 32, 32, 512)  2359808     ['dropout_4[0][0]']              \n",
      "                                                                                                  \n",
      " activation_8 (Activation)      (None, 32, 32, 512)  0           ['middle_conv1-2[0][0]']         \n",
      "                                                                                                  \n",
      " dropout_5 (Dropout)            (None, 32, 32, 512)  0           ['activation_8[0][0]']           \n",
      "                                                                                                  \n",
      " upsampling_4 (Conv2DTranspose)  (None, 64, 64, 256)  1179904    ['dropout_5[0][0]']              \n",
      "                                                                                                  \n",
      " concatenate (Concatenate)      (None, 64, 64, 512)  0           ['upsampling_4[0][0]',           \n",
      "                                                                  'activation_6[0][0]']           \n",
      "                                                                                                  \n",
      " up_conv4-1 (Conv2D)            (None, 64, 64, 256)  1179904     ['concatenate[0][0]']            \n",
      "                                                                                                  \n",
      " activation_9 (Activation)      (None, 64, 64, 256)  0           ['up_conv4-1[0][0]']             \n",
      "                                                                                                  \n",
      " dropout_6 (Dropout)            (None, 64, 64, 256)  0           ['activation_9[0][0]']           \n",
      "                                                                                                  \n",
      " up_conv4-2 (Conv2D)            (None, 64, 64, 256)  590080      ['dropout_6[0][0]']              \n",
      "                                                                                                  \n",
      " activation_10 (Activation)     (None, 64, 64, 256)  0           ['up_conv4-2[0][0]']             \n",
      "                                                                                                  \n",
      " dropout_7 (Dropout)            (None, 64, 64, 256)  0           ['activation_10[0][0]']          \n",
      "                                                                                                  \n",
      " upsampling_3 (Conv2DTranspose)  (None, 128, 128, 12  295040     ['dropout_7[0][0]']              \n",
      "                                8)                                                                \n",
      "                                                                                                  \n",
      " concatenate_1 (Concatenate)    (None, 128, 128, 25  0           ['upsampling_3[0][0]',           \n",
      "                                6)                                'activation_4[0][0]']           \n",
      "                                                                                                  \n",
      " up_conv3-1 (Conv2D)            (None, 128, 128, 12  295040      ['concatenate_1[0][0]']          \n",
      "                                8)                                                                \n",
      "                                                                                                  \n",
      " activation_11 (Activation)     (None, 128, 128, 12  0           ['up_conv3-1[0][0]']             \n",
      "                                8)                                                                \n",
      "                                                                                                  \n",
      " dropout_8 (Dropout)            (None, 128, 128, 12  0           ['activation_11[0][0]']          \n",
      "                                8)                                                                \n",
      "                                                                                                  \n",
      " up_conv3-2 (Conv2D)            (None, 128, 128, 12  147584      ['dropout_8[0][0]']              \n",
      "                                8)                                                                \n",
      "                                                                                                  \n",
      " activation_12 (Activation)     (None, 128, 128, 12  0           ['up_conv3-2[0][0]']             \n",
      "                                8)                                                                \n",
      "                                                                                                  \n",
      " dropout_9 (Dropout)            (None, 128, 128, 12  0           ['activation_12[0][0]']          \n",
      "                                8)                                                                \n",
      "                                                                                                  \n",
      " upsampling_2 (Conv2DTranspose)  (None, 256, 256, 64  73792      ['dropout_9[0][0]']              \n",
      "                                )                                                                 \n",
      "                                                                                                  \n",
      " concatenate_2 (Concatenate)    (None, 256, 256, 12  0           ['upsampling_2[0][0]',           \n",
      "                                8)                                'activation_2[0][0]']           \n",
      "                                                                                                  \n",
      " up_conv2-1 (Conv2D)            (None, 256, 256, 64  73792       ['concatenate_2[0][0]']          \n",
      "                                )                                                                 \n",
      "                                                                                                  \n",
      " activation_13 (Activation)     (None, 256, 256, 64  0           ['up_conv2-1[0][0]']             \n",
      "                                )                                                                 \n",
      "                                                                                                  \n",
      " dropout_10 (Dropout)           (None, 256, 256, 64  0           ['activation_13[0][0]']          \n",
      "                                )                                                                 \n",
      "                                                                                                  \n",
      " up_conv2-2 (Conv2D)            (None, 256, 256, 64  36928       ['dropout_10[0][0]']             \n",
      "                                )                                                                 \n",
      "                                                                                                  \n",
      " activation_14 (Activation)     (None, 256, 256, 64  0           ['up_conv2-2[0][0]']             \n",
      "                                )                                                                 \n",
      "                                                                                                  \n",
      " dropout_11 (Dropout)           (None, 256, 256, 64  0           ['activation_14[0][0]']          \n",
      "                                )                                                                 \n",
      "                                                                                                  \n",
      " upsampling_1 (Conv2DTranspose)  (None, 512, 512, 32  18464      ['dropout_11[0][0]']             \n",
      "                                )                                                                 \n",
      "                                                                                                  \n",
      " concatenate_3 (Concatenate)    (None, 512, 512, 64  0           ['upsampling_1[0][0]',           \n",
      "                                )                                 'activation[0][0]']             \n",
      "                                                                                                  \n",
      " up_conv1 (Conv2D)              (None, 512, 512, 32  18464       ['concatenate_3[0][0]']          \n",
      "                                )                                                                 \n",
      "                                                                                                  \n",
      " activation_15 (Activation)     (None, 512, 512, 32  0           ['up_conv1[0][0]']               \n",
      "                                )                                                                 \n",
      "                                                                                                  \n",
      " dropout_12 (Dropout)           (None, 512, 512, 32  0           ['activation_15[0][0]']          \n",
      "                                )                                                                 \n",
      "                                                                                                  \n",
      " segmentation_layer (Conv2D)    (None, 512, 512, 1)  33          ['dropout_12[0][0]']             \n",
      "                                                                                                  \n",
      "==================================================================================================\n",
      "Total params: 8,611,425\n",
      "Trainable params: 8,611,425\n",
      "Non-trainable params: 0\n",
      "__________________________________________________________________________________________________\n"
     ]
    }
   ],
   "source": [
    "model.summary()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e7144803-17e2-4947-98d2-e0db6bdc8102",
   "metadata": {
    "jp-MarkdownHeadingCollapsed": true,
    "tags": []
   },
   "source": [
    "## Test AM-UNet 2D"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "56883bea-78c0-42a0-b903-983782777438",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "2254065\n",
      "{'current': '0.370', 'peak': '0.371'}\n"
     ]
    }
   ],
   "source": [
    "importlib.reload(unet_arch)\n",
    "\n",
    "model = unet_arch.UNet(\n",
    "    input_shape = (512, 512, 1),\n",
    "    output_dim  = 1,\n",
    "    \n",
    "    n_stages    = 5,\n",
    "    n_conv_per_stage = lambda i: 1 if i <= 1 else 2,\n",
    "    filters     = [16, 32, 64, 128, 128],\n",
    "    bnorm       = 'after',\n",
    "    activation  = 'relu',\n",
    "    drop_rate   = lambda i: 0. if i == 0 else 0.25,\n",
    "    \n",
    "    n_middle_stages = 4,\n",
    "    n_middle_conv   = 2,\n",
    "    middle_filters  = 64,\n",
    "    middle_bnorm    = 'after'\n",
    ")\n",
    "print(model.count_params())\n",
    "model.compile(loss = 'binary_crossentropy', optimizer = 'adam')\n",
    "\n",
    "show_memory()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "e5d49d54-91d2-4e15-8f26-03596e4848e3",
   "metadata": {
    "collapsed": true,
    "jupyter": {
     "outputs_hidden": true
    },
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Before\t: {'current': '0.298', 'peak': '0.298'}\n",
      "Epoch 1/5\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "2023-03-23 16:02:17.752060: W tensorflow/core/kernels/gpu_utils.cc:50] Failed to allocate memory for convolution redzone checking; skipping this check. This is benign and only means that we won't check cudnn for out-of-bounds reads and writes. This message will only be printed once.\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "2/2 [==============================] - 10s 719ms/step - loss: 0.4668\n",
      "Epoch 2/5\n",
      "2/2 [==============================] - 1s 718ms/step - loss: 0.4478\n",
      "Epoch 3/5\n",
      "2/2 [==============================] - 1s 721ms/step - loss: 0.4291\n",
      "Epoch 4/5\n",
      "2/2 [==============================] - 1s 720ms/step - loss: 0.4109\n",
      "Epoch 5/5\n",
      "2/2 [==============================] - 1s 720ms/step - loss: 0.3937\n",
      "After fit\t: {'current': '0.484', 'peak': '13.884'}\n"
     ]
    }
   ],
   "source": [
    "batch_size = 42\n",
    "\n",
    "show_memory('Before')\n",
    "\n",
    "inp = tf.random.uniform((batch_size * 2, ) + tuple(model.input_shape[1:]), 0., 1.)\n",
    "out = tf.ones_like(inp)\n",
    "\n",
    "\"\"\"show_memory('Batch init')\n",
    "\n",
    "model(inp[:batch_size], training = True)\n",
    "\n",
    "show_memory('Simple call')\n",
    "\n",
    "with tf.GradientTape() as tape:\n",
    "    pred = model(inp[:batch_size], training = True)\n",
    "    l = model.compiled_loss(out[:batch_size], pred)\n",
    "\n",
    "show_memory('With gradient')\n",
    "\n",
    "del tape, l, pred\n",
    "\n",
    "show_memory('After deleting tape')\n",
    "\"\"\"\n",
    "_ = model.fit(inp, out, epochs = 5, batch_size = batch_size)\n",
    "\n",
    "show_memory('After fit')\n",
    "tf.config.experimental.reset_memory_stats('GPU:0')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "5a67149a-b819-467a-bd7c-a9ea52749a99",
   "metadata": {
    "collapsed": true,
    "jupyter": {
     "outputs_hidden": true
    },
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Model: \"model\"\n",
      "__________________________________________________________________________________________________\n",
      " Layer (type)                   Output Shape         Param #     Connected to                     \n",
      "==================================================================================================\n",
      " input_image (InputLayer)       [(None, 512, 512, 1  0           []                               \n",
      "                                )]                                                                \n",
      "                                                                                                  \n",
      " down_conv1 (Conv2D)            (None, 512, 512, 16  160         ['input_image[0][0]']            \n",
      "                                )                                                                 \n",
      "                                                                                                  \n",
      " down_bn1 (BatchNormalization)  (None, 512, 512, 16  64          ['down_conv1[0][0]']             \n",
      "                                )                                                                 \n",
      "                                                                                                  \n",
      " activation (Activation)        (None, 512, 512, 16  0           ['down_bn1[0][0]']               \n",
      "                                )                                                                 \n",
      "                                                                                                  \n",
      " max_pooling2d (MaxPooling2D)   (None, 256, 256, 16  0           ['activation[0][0]']             \n",
      "                                )                                                                 \n",
      "                                                                                                  \n",
      " down_conv2-1 (Conv2D)          (None, 256, 256, 32  4640        ['max_pooling2d[0][0]']          \n",
      "                                )                                                                 \n",
      "                                                                                                  \n",
      " down_bn2-1 (BatchNormalization  (None, 256, 256, 32  128        ['down_conv2-1[0][0]']           \n",
      " )                              )                                                                 \n",
      "                                                                                                  \n",
      " activation_1 (Activation)      (None, 256, 256, 32  0           ['down_bn2-1[0][0]']             \n",
      "                                )                                                                 \n",
      "                                                                                                  \n",
      " down_conv2-2 (Conv2D)          (None, 256, 256, 32  9248        ['activation_1[0][0]']           \n",
      "                                )                                                                 \n",
      "                                                                                                  \n",
      " down_bn2-2 (BatchNormalization  (None, 256, 256, 32  128        ['down_conv2-2[0][0]']           \n",
      " )                              )                                                                 \n",
      "                                                                                                  \n",
      " activation_2 (Activation)      (None, 256, 256, 32  0           ['down_bn2-2[0][0]']             \n",
      "                                )                                                                 \n",
      "                                                                                                  \n",
      " max_pooling2d_1 (MaxPooling2D)  (None, 128, 128, 32  0          ['activation_2[0][0]']           \n",
      "                                )                                                                 \n",
      "                                                                                                  \n",
      " dropout (Dropout)              (None, 128, 128, 32  0           ['max_pooling2d_1[0][0]']        \n",
      "                                )                                                                 \n",
      "                                                                                                  \n",
      " down_conv3-1 (Conv2D)          (None, 128, 128, 64  18496       ['dropout[0][0]']                \n",
      "                                )                                                                 \n",
      "                                                                                                  \n",
      " down_bn3-1 (BatchNormalization  (None, 128, 128, 64  256        ['down_conv3-1[0][0]']           \n",
      " )                              )                                                                 \n",
      "                                                                                                  \n",
      " activation_3 (Activation)      (None, 128, 128, 64  0           ['down_bn3-1[0][0]']             \n",
      "                                )                                                                 \n",
      "                                                                                                  \n",
      " down_conv3-2 (Conv2D)          (None, 128, 128, 64  36928       ['activation_3[0][0]']           \n",
      "                                )                                                                 \n",
      "                                                                                                  \n",
      " down_bn3-2 (BatchNormalization  (None, 128, 128, 64  256        ['down_conv3-2[0][0]']           \n",
      " )                              )                                                                 \n",
      "                                                                                                  \n",
      " activation_4 (Activation)      (None, 128, 128, 64  0           ['down_bn3-2[0][0]']             \n",
      "                                )                                                                 \n",
      "                                                                                                  \n",
      " max_pooling2d_2 (MaxPooling2D)  (None, 64, 64, 64)  0           ['activation_4[0][0]']           \n",
      "                                                                                                  \n",
      " dropout_1 (Dropout)            (None, 64, 64, 64)   0           ['max_pooling2d_2[0][0]']        \n",
      "                                                                                                  \n",
      " down_conv4-1 (Conv2D)          (None, 64, 64, 128)  73856       ['dropout_1[0][0]']              \n",
      "                                                                                                  \n",
      " down_bn4-1 (BatchNormalization  (None, 64, 64, 128)  512        ['down_conv4-1[0][0]']           \n",
      " )                                                                                                \n",
      "                                                                                                  \n",
      " activation_5 (Activation)      (None, 64, 64, 128)  0           ['down_bn4-1[0][0]']             \n",
      "                                                                                                  \n",
      " down_conv4-2 (Conv2D)          (None, 64, 64, 128)  147584      ['activation_5[0][0]']           \n",
      "                                                                                                  \n",
      " down_bn4-2 (BatchNormalization  (None, 64, 64, 128)  512        ['down_conv4-2[0][0]']           \n",
      " )                                                                                                \n",
      "                                                                                                  \n",
      " activation_6 (Activation)      (None, 64, 64, 128)  0           ['down_bn4-2[0][0]']             \n",
      "                                                                                                  \n",
      " max_pooling2d_3 (MaxPooling2D)  (None, 32, 32, 128)  0          ['activation_6[0][0]']           \n",
      "                                                                                                  \n",
      " dropout_2 (Dropout)            (None, 32, 32, 128)  0           ['max_pooling2d_3[0][0]']        \n",
      "                                                                                                  \n",
      " down_conv5-1 (Conv2D)          (None, 32, 32, 128)  147584      ['dropout_2[0][0]']              \n",
      "                                                                                                  \n",
      " down_bn5-1 (BatchNormalization  (None, 32, 32, 128)  512        ['down_conv5-1[0][0]']           \n",
      " )                                                                                                \n",
      "                                                                                                  \n",
      " activation_7 (Activation)      (None, 32, 32, 128)  0           ['down_bn5-1[0][0]']             \n",
      "                                                                                                  \n",
      " down_conv5-2 (Conv2D)          (None, 32, 32, 128)  147584      ['activation_7[0][0]']           \n",
      "                                                                                                  \n",
      " down_bn5-2 (BatchNormalization  (None, 32, 32, 128)  512        ['down_conv5-2[0][0]']           \n",
      " )                                                                                                \n",
      "                                                                                                  \n",
      " activation_8 (Activation)      (None, 32, 32, 128)  0           ['down_bn5-2[0][0]']             \n",
      "                                                                                                  \n",
      " max_pooling2d_4 (MaxPooling2D)  (None, 16, 16, 128)  0          ['activation_8[0][0]']           \n",
      "                                                                                                  \n",
      " dropout_3 (Dropout)            (None, 16, 16, 128)  0           ['max_pooling2d_4[0][0]']        \n",
      "                                                                                                  \n",
      " middle_conv1-1 (Conv2D)        (None, 16, 16, 64)   73792       ['dropout_3[0][0]']              \n",
      "                                                                                                  \n",
      " middle_bn1-1 (BatchNormalizati  (None, 16, 16, 64)  256         ['middle_conv1-1[0][0]']         \n",
      " on)                                                                                              \n",
      "                                                                                                  \n",
      " activation_9 (Activation)      (None, 16, 16, 64)   0           ['middle_bn1-1[0][0]']           \n",
      "                                                                                                  \n",
      " dropout_4 (Dropout)            (None, 16, 16, 64)   0           ['activation_9[0][0]']           \n",
      "                                                                                                  \n",
      " middle_conv1-2 (Conv2D)        (None, 16, 16, 64)   36928       ['dropout_4[0][0]']              \n",
      "                                                                                                  \n",
      " middle_bn1-2 (BatchNormalizati  (None, 16, 16, 64)  256         ['middle_conv1-2[0][0]']         \n",
      " on)                                                                                              \n",
      "                                                                                                  \n",
      " activation_10 (Activation)     (None, 16, 16, 64)   0           ['middle_bn1-2[0][0]']           \n",
      "                                                                                                  \n",
      " dropout_5 (Dropout)            (None, 16, 16, 64)   0           ['activation_10[0][0]']          \n",
      "                                                                                                  \n",
      " middle_conv2-1 (Conv2D)        (None, 16, 16, 64)   36928       ['dropout_5[0][0]']              \n",
      "                                                                                                  \n",
      " middle_bn2-1 (BatchNormalizati  (None, 16, 16, 64)  256         ['middle_conv2-1[0][0]']         \n",
      " on)                                                                                              \n",
      "                                                                                                  \n",
      " activation_11 (Activation)     (None, 16, 16, 64)   0           ['middle_bn2-1[0][0]']           \n",
      "                                                                                                  \n",
      " dropout_6 (Dropout)            (None, 16, 16, 64)   0           ['activation_11[0][0]']          \n",
      "                                                                                                  \n",
      " middle_conv2-2 (Conv2D)        (None, 16, 16, 64)   36928       ['dropout_6[0][0]']              \n",
      "                                                                                                  \n",
      " middle_bn2-2 (BatchNormalizati  (None, 16, 16, 64)  256         ['middle_conv2-2[0][0]']         \n",
      " on)                                                                                              \n",
      "                                                                                                  \n",
      " activation_12 (Activation)     (None, 16, 16, 64)   0           ['middle_bn2-2[0][0]']           \n",
      "                                                                                                  \n",
      " dropout_7 (Dropout)            (None, 16, 16, 64)   0           ['activation_12[0][0]']          \n",
      "                                                                                                  \n",
      " middle_conv3-1 (Conv2D)        (None, 16, 16, 64)   36928       ['dropout_7[0][0]']              \n",
      "                                                                                                  \n",
      " middle_bn3-1 (BatchNormalizati  (None, 16, 16, 64)  256         ['middle_conv3-1[0][0]']         \n",
      " on)                                                                                              \n",
      "                                                                                                  \n",
      " activation_13 (Activation)     (None, 16, 16, 64)   0           ['middle_bn3-1[0][0]']           \n",
      "                                                                                                  \n",
      " dropout_8 (Dropout)            (None, 16, 16, 64)   0           ['activation_13[0][0]']          \n",
      "                                                                                                  \n",
      " middle_conv3-2 (Conv2D)        (None, 16, 16, 64)   36928       ['dropout_8[0][0]']              \n",
      "                                                                                                  \n",
      " middle_bn3-2 (BatchNormalizati  (None, 16, 16, 64)  256         ['middle_conv3-2[0][0]']         \n",
      " on)                                                                                              \n",
      "                                                                                                  \n",
      " activation_14 (Activation)     (None, 16, 16, 64)   0           ['middle_bn3-2[0][0]']           \n",
      "                                                                                                  \n",
      " dropout_9 (Dropout)            (None, 16, 16, 64)   0           ['activation_14[0][0]']          \n",
      "                                                                                                  \n",
      " middle_conv4-1 (Conv2D)        (None, 16, 16, 64)   36928       ['dropout_9[0][0]']              \n",
      "                                                                                                  \n",
      " middle_bn4-1 (BatchNormalizati  (None, 16, 16, 64)  256         ['middle_conv4-1[0][0]']         \n",
      " on)                                                                                              \n",
      "                                                                                                  \n",
      " activation_15 (Activation)     (None, 16, 16, 64)   0           ['middle_bn4-1[0][0]']           \n",
      "                                                                                                  \n",
      " dropout_10 (Dropout)           (None, 16, 16, 64)   0           ['activation_15[0][0]']          \n",
      "                                                                                                  \n",
      " middle_conv4-2 (Conv2D)        (None, 16, 16, 64)   36928       ['dropout_10[0][0]']             \n",
      "                                                                                                  \n",
      " middle_bn4-2 (BatchNormalizati  (None, 16, 16, 64)  256         ['middle_conv4-2[0][0]']         \n",
      " on)                                                                                              \n",
      "                                                                                                  \n",
      " activation_16 (Activation)     (None, 16, 16, 64)   0           ['middle_bn4-2[0][0]']           \n",
      "                                                                                                  \n",
      " dropout_11 (Dropout)           (None, 16, 16, 64)   0           ['activation_16[0][0]']          \n",
      "                                                                                                  \n",
      " upsampling_5 (Conv2DTranspose)  (None, 32, 32, 128)  73856      ['dropout_11[0][0]']             \n",
      "                                                                                                  \n",
      " concatenate (Concatenate)      (None, 32, 32, 256)  0           ['upsampling_5[0][0]',           \n",
      "                                                                  'activation_8[0][0]']           \n",
      "                                                                                                  \n",
      " up_conv5-1 (Conv2D)            (None, 32, 32, 128)  295040      ['concatenate[0][0]']            \n",
      "                                                                                                  \n",
      " up_bn5-1 (BatchNormalization)  (None, 32, 32, 128)  512         ['up_conv5-1[0][0]']             \n",
      "                                                                                                  \n",
      " activation_17 (Activation)     (None, 32, 32, 128)  0           ['up_bn5-1[0][0]']               \n",
      "                                                                                                  \n",
      " up_conv5-2 (Conv2D)            (None, 32, 32, 128)  147584      ['activation_17[0][0]']          \n",
      "                                                                                                  \n",
      " up_bn5-2 (BatchNormalization)  (None, 32, 32, 128)  512         ['up_conv5-2[0][0]']             \n",
      "                                                                                                  \n",
      " activation_18 (Activation)     (None, 32, 32, 128)  0           ['up_bn5-2[0][0]']               \n",
      "                                                                                                  \n",
      " dropout_12 (Dropout)           (None, 32, 32, 128)  0           ['activation_18[0][0]']          \n",
      "                                                                                                  \n",
      " upsampling_4 (Conv2DTranspose)  (None, 64, 64, 128)  147584     ['dropout_12[0][0]']             \n",
      "                                                                                                  \n",
      " concatenate_1 (Concatenate)    (None, 64, 64, 256)  0           ['upsampling_4[0][0]',           \n",
      "                                                                  'activation_6[0][0]']           \n",
      "                                                                                                  \n",
      " up_conv4-1 (Conv2D)            (None, 64, 64, 128)  295040      ['concatenate_1[0][0]']          \n",
      "                                                                                                  \n",
      " up_bn4-1 (BatchNormalization)  (None, 64, 64, 128)  512         ['up_conv4-1[0][0]']             \n",
      "                                                                                                  \n",
      " activation_19 (Activation)     (None, 64, 64, 128)  0           ['up_bn4-1[0][0]']               \n",
      "                                                                                                  \n",
      " up_conv4-2 (Conv2D)            (None, 64, 64, 128)  147584      ['activation_19[0][0]']          \n",
      "                                                                                                  \n",
      " up_bn4-2 (BatchNormalization)  (None, 64, 64, 128)  512         ['up_conv4-2[0][0]']             \n",
      "                                                                                                  \n",
      " activation_20 (Activation)     (None, 64, 64, 128)  0           ['up_bn4-2[0][0]']               \n",
      "                                                                                                  \n",
      " dropout_13 (Dropout)           (None, 64, 64, 128)  0           ['activation_20[0][0]']          \n",
      "                                                                                                  \n",
      " upsampling_3 (Conv2DTranspose)  (None, 128, 128, 64  73792      ['dropout_13[0][0]']             \n",
      "                                )                                                                 \n",
      "                                                                                                  \n",
      " concatenate_2 (Concatenate)    (None, 128, 128, 12  0           ['upsampling_3[0][0]',           \n",
      "                                8)                                'activation_4[0][0]']           \n",
      "                                                                                                  \n",
      " up_conv3-1 (Conv2D)            (None, 128, 128, 64  73792       ['concatenate_2[0][0]']          \n",
      "                                )                                                                 \n",
      "                                                                                                  \n",
      " up_bn3-1 (BatchNormalization)  (None, 128, 128, 64  256         ['up_conv3-1[0][0]']             \n",
      "                                )                                                                 \n",
      "                                                                                                  \n",
      " activation_21 (Activation)     (None, 128, 128, 64  0           ['up_bn3-1[0][0]']               \n",
      "                                )                                                                 \n",
      "                                                                                                  \n",
      " up_conv3-2 (Conv2D)            (None, 128, 128, 64  36928       ['activation_21[0][0]']          \n",
      "                                )                                                                 \n",
      "                                                                                                  \n",
      " up_bn3-2 (BatchNormalization)  (None, 128, 128, 64  256         ['up_conv3-2[0][0]']             \n",
      "                                )                                                                 \n",
      "                                                                                                  \n",
      " activation_22 (Activation)     (None, 128, 128, 64  0           ['up_bn3-2[0][0]']               \n",
      "                                )                                                                 \n",
      "                                                                                                  \n",
      " dropout_14 (Dropout)           (None, 128, 128, 64  0           ['activation_22[0][0]']          \n",
      "                                )                                                                 \n",
      "                                                                                                  \n",
      " upsampling_2 (Conv2DTranspose)  (None, 256, 256, 32  18464      ['dropout_14[0][0]']             \n",
      "                                )                                                                 \n",
      "                                                                                                  \n",
      " concatenate_3 (Concatenate)    (None, 256, 256, 64  0           ['upsampling_2[0][0]',           \n",
      "                                )                                 'activation_2[0][0]']           \n",
      "                                                                                                  \n",
      " up_conv2-1 (Conv2D)            (None, 256, 256, 32  18464       ['concatenate_3[0][0]']          \n",
      "                                )                                                                 \n",
      "                                                                                                  \n",
      " up_bn2-1 (BatchNormalization)  (None, 256, 256, 32  128         ['up_conv2-1[0][0]']             \n",
      "                                )                                                                 \n",
      "                                                                                                  \n",
      " activation_23 (Activation)     (None, 256, 256, 32  0           ['up_bn2-1[0][0]']               \n",
      "                                )                                                                 \n",
      "                                                                                                  \n",
      " up_conv2-2 (Conv2D)            (None, 256, 256, 32  9248        ['activation_23[0][0]']          \n",
      "                                )                                                                 \n",
      "                                                                                                  \n",
      " up_bn2-2 (BatchNormalization)  (None, 256, 256, 32  128         ['up_conv2-2[0][0]']             \n",
      "                                )                                                                 \n",
      "                                                                                                  \n",
      " activation_24 (Activation)     (None, 256, 256, 32  0           ['up_bn2-2[0][0]']               \n",
      "                                )                                                                 \n",
      "                                                                                                  \n",
      " dropout_15 (Dropout)           (None, 256, 256, 32  0           ['activation_24[0][0]']          \n",
      "                                )                                                                 \n",
      "                                                                                                  \n",
      " upsampling_1 (Conv2DTranspose)  (None, 512, 512, 16  4624       ['dropout_15[0][0]']             \n",
      "                                )                                                                 \n",
      "                                                                                                  \n",
      " concatenate_4 (Concatenate)    (None, 512, 512, 32  0           ['upsampling_1[0][0]',           \n",
      "                                )                                 'activation[0][0]']             \n",
      "                                                                                                  \n",
      " up_conv1 (Conv2D)              (None, 512, 512, 16  4624        ['concatenate_4[0][0]']          \n",
      "                                )                                                                 \n",
      "                                                                                                  \n",
      " up_bn1 (BatchNormalization)    (None, 512, 512, 16  64          ['up_conv1[0][0]']               \n",
      "                                )                                                                 \n",
      "                                                                                                  \n",
      " activation_25 (Activation)     (None, 512, 512, 16  0           ['up_bn1[0][0]']                 \n",
      "                                )                                                                 \n",
      "                                                                                                  \n",
      " segmentation_layer (Conv2D)    (None, 512, 512, 1)  17          ['activation_25[0][0]']          \n",
      "                                                                                                  \n",
      "==================================================================================================\n",
      "Total params: 2,272,817\n",
      "Trainable params: 2,268,913\n",
      "Non-trainable params: 3,904\n",
      "__________________________________________________________________________________________________\n"
     ]
    }
   ],
   "source": [
    "model.summary()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "77a037bc-72b1-4d69-9862-e97e96c835c9",
   "metadata": {
    "jp-MarkdownHeadingCollapsed": true,
    "tags": []
   },
   "source": [
    "## Test AM-UNet 3D"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "e0f27eab-e049-4556-b2fc-6d6ece8de5eb",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "INFO:tensorflow:Mixed precision compatibility check (mixed_float16): OK\n",
      "Your GPUs will likely run quickly with dtype policy mixed_float16 as they all have compute capability of at least 7.0\n",
      "Mixed precision compatibility check (mixed_float16): OK\n",
      "Your GPUs will likely run quickly with dtype policy mixed_float16 as they all have compute capability of at least 7.0\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "2023-03-24 10:26:09.319656: I tensorflow/core/platform/cpu_feature_guard.cc:193] This TensorFlow binary is optimized with oneAPI Deep Neural Network Library (oneDNN) to use the following CPU instructions in performance-critical operations:  AVX2 AVX512F AVX512_VNNI FMA\n",
      "To enable them in other operations, rebuild TensorFlow with the appropriate compiler flags.\n",
      "2023-03-24 10:26:09.691052: I tensorflow/core/common_runtime/gpu/gpu_device.cc:1616] Created device /job:localhost/replica:0/task:0/device:GPU:0 with 14783 MB memory:  -> device: 0, name: Quadro RTX 5000, pci bus id: 0000:17:00.0, compute capability: 7.5\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "3416209\n",
      "{'current': '0.013', 'peak': '0.016'}\n"
     ]
    }
   ],
   "source": [
    "importlib.reload(unet_arch)\n",
    "\n",
    "model = unet_arch.UNet(\n",
    "    input_shape = (None, 512, 512, 1),\n",
    "    output_dim  = 1,\n",
    "    \n",
    "    n_stages    = 4,\n",
    "    n_conv_per_stage = lambda i: 1 if i <= 2 else 2,\n",
    "    filters     = [16, 32, 64, 128, 128],\n",
    "    bnorm       = 'never',\n",
    "    activation  = 'relu',\n",
    "    drop_rate   = lambda i: 0. if i == 0 else 0.25,\n",
    "    \n",
    "    n_middle_stages = 2,\n",
    "    n_middle_conv   = 2,\n",
    "    middle_filters  = 64,\n",
    "    middle_bnorm    = 'never',\n",
    "    \n",
    "    mixed_precision = True\n",
    ")\n",
    "print(model.count_params())\n",
    "model.compile(loss = 'binary_crossentropy', optimizer = 'adam')\n",
    "\n",
    "show_memory()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "acaba06f-055e-434e-84cb-d5f52087dd8b",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Before\t: {'current': '0.518', 'peak': '7.675'}\n",
      "Epoch 1/5\n",
      "2/2 [==============================] - 3s 2s/step - loss: 0.0000e+00\n",
      "Epoch 2/5\n",
      "2/2 [==============================] - 3s 2s/step - loss: 0.0000e+00\n",
      "Epoch 3/5\n",
      "2/2 [==============================] - 3s 2s/step - loss: 0.0000e+00\n",
      "Epoch 4/5\n",
      "2/2 [==============================] - 3s 2s/step - loss: 0.0000e+00\n",
      "Epoch 5/5\n",
      "2/2 [==============================] - 3s 2s/step - loss: 0.0000e+00\n",
      "After fit\t: {'current': '0.768', 'peak': '12.750'}\n"
     ]
    }
   ],
   "source": [
    "batch_size = 1\n",
    "seq_len    = 64\n",
    "\n",
    "show_memory('Before')\n",
    "\n",
    "inp = tf.random.uniform((batch_size * 2, seq_len) + tuple(model.input_shape[2:]), 0., 1.)\n",
    "out = tf.ones_like(inp)\n",
    "\n",
    "_ = model.fit(inp, out, epochs = 5, batch_size = batch_size)\n",
    "\n",
    "show_memory('After fit')\n",
    "tf.config.experimental.reset_memory_stats('GPU:0')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "b36cd952-c214-4fda-9040-89579725e129",
   "metadata": {
    "collapsed": true,
    "jupyter": {
     "outputs_hidden": true
    },
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "64\n",
      "Before\t: {'current': '0.013', 'peak': '0.016'}\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "2023-03-24 10:26:16.412280: I tensorflow/stream_executor/cuda/cuda_dnn.cc:384] Loaded cuDNN version 8204\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "After fit\t: {'current': '0.200', 'peak': '3.888'}\n",
      "96\n",
      "Before\t: {'current': '0.200', 'peak': '0.200'}\n",
      "After fit\t: {'current': '0.294', 'peak': '5.935'}\n",
      "128\n",
      "Before\t: {'current': '0.294', 'peak': '0.294'}\n",
      "After fit\t: {'current': '0.450', 'peak': '7.920'}\n",
      "160\n",
      "Before\t: {'current': '0.450', 'peak': '0.450'}\n",
      "After fit\t: {'current': '0.575', 'peak': '9.967'}\n",
      "192\n",
      "Before\t: {'current': '0.575', 'peak': '0.575'}\n",
      "After fit\t: {'current': '0.638', 'peak': '11.983'}\n",
      "224\n",
      "Before\t: {'current': '0.638', 'peak': '0.638'}\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "2023-03-24 10:26:36.306704: W tensorflow/core/common_runtime/bfc_allocator.cc:290] Allocator (GPU_0_bfc) ran out of memory trying to allocate 6.25GiB with freed_by_count=0. The caller indicates that this is not a failure, but this may mean that there could be performance gains if more memory were available.\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "After fit\t: {'current': '0.669', 'peak': '10.369'}\n",
      "256\n",
      "Before\t: {'current': '0.669', 'peak': '0.669'}\n"
     ]
    },
    {
     "ename": "KeyboardInterrupt",
     "evalue": "",
     "output_type": "error",
     "traceback": [
      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[0;31mKeyboardInterrupt\u001b[0m                         Traceback (most recent call last)",
      "Cell \u001b[0;32mIn [3], line 8\u001b[0m\n\u001b[1;32m      5\u001b[0m show_memory(\u001b[38;5;124m'\u001b[39m\u001b[38;5;124mBefore\u001b[39m\u001b[38;5;124m'\u001b[39m)\n\u001b[1;32m      7\u001b[0m inp \u001b[38;5;241m=\u001b[39m tf\u001b[38;5;241m.\u001b[39mrandom\u001b[38;5;241m.\u001b[39muniform((batch_size, seq_len) \u001b[38;5;241m+\u001b[39m \u001b[38;5;28mtuple\u001b[39m(model\u001b[38;5;241m.\u001b[39minput_shape[\u001b[38;5;241m2\u001b[39m:]), \u001b[38;5;241m0.\u001b[39m, \u001b[38;5;241m1.\u001b[39m)\n\u001b[0;32m----> 8\u001b[0m _ \u001b[38;5;241m=\u001b[39m \u001b[43mmodel\u001b[49m\u001b[43m(\u001b[49m\u001b[43minp\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m     10\u001b[0m show_memory(\u001b[38;5;124m'\u001b[39m\u001b[38;5;124mAfter fit\u001b[39m\u001b[38;5;124m'\u001b[39m)\n\u001b[1;32m     11\u001b[0m tf\u001b[38;5;241m.\u001b[39mconfig\u001b[38;5;241m.\u001b[39mexperimental\u001b[38;5;241m.\u001b[39mreset_memory_stats(\u001b[38;5;124m'\u001b[39m\u001b[38;5;124mGPU:0\u001b[39m\u001b[38;5;124m'\u001b[39m)\n",
      "File \u001b[0;32m~/.local/lib/python3.9/site-packages/keras/utils/traceback_utils.py:65\u001b[0m, in \u001b[0;36mfilter_traceback.<locals>.error_handler\u001b[0;34m(*args, **kwargs)\u001b[0m\n\u001b[1;32m     63\u001b[0m filtered_tb \u001b[38;5;241m=\u001b[39m \u001b[38;5;28;01mNone\u001b[39;00m\n\u001b[1;32m     64\u001b[0m \u001b[38;5;28;01mtry\u001b[39;00m:\n\u001b[0;32m---> 65\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[43mfn\u001b[49m\u001b[43m(\u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43margs\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43mkwargs\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m     66\u001b[0m \u001b[38;5;28;01mexcept\u001b[39;00m \u001b[38;5;167;01mException\u001b[39;00m \u001b[38;5;28;01mas\u001b[39;00m e:\n\u001b[1;32m     67\u001b[0m     filtered_tb \u001b[38;5;241m=\u001b[39m _process_traceback_frames(e\u001b[38;5;241m.\u001b[39m__traceback__)\n",
      "File \u001b[0;32m~/.local/lib/python3.9/site-packages/keras/engine/training.py:557\u001b[0m, in \u001b[0;36mModel.__call__\u001b[0;34m(self, *args, **kwargs)\u001b[0m\n\u001b[1;32m    553\u001b[0m         \u001b[38;5;28msuper\u001b[39m()\u001b[38;5;241m.\u001b[39m\u001b[38;5;21m__call__\u001b[39m(inputs, \u001b[38;5;241m*\u001b[39mcopied_args, \u001b[38;5;241m*\u001b[39m\u001b[38;5;241m*\u001b[39mcopied_kwargs)\n\u001b[1;32m    555\u001b[0m     layout_map_lib\u001b[38;5;241m.\u001b[39m_map_subclass_model_variable(\u001b[38;5;28mself\u001b[39m, \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_layout_map)\n\u001b[0;32m--> 557\u001b[0m \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[38;5;28;43msuper\u001b[39;49m\u001b[43m(\u001b[49m\u001b[43m)\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[38;5;21;43m__call__\u001b[39;49m\u001b[43m(\u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43margs\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43mkwargs\u001b[49m\u001b[43m)\u001b[49m\n",
      "File \u001b[0;32m~/.local/lib/python3.9/site-packages/keras/utils/traceback_utils.py:65\u001b[0m, in \u001b[0;36mfilter_traceback.<locals>.error_handler\u001b[0;34m(*args, **kwargs)\u001b[0m\n\u001b[1;32m     63\u001b[0m filtered_tb \u001b[38;5;241m=\u001b[39m \u001b[38;5;28;01mNone\u001b[39;00m\n\u001b[1;32m     64\u001b[0m \u001b[38;5;28;01mtry\u001b[39;00m:\n\u001b[0;32m---> 65\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[43mfn\u001b[49m\u001b[43m(\u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43margs\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43mkwargs\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m     66\u001b[0m \u001b[38;5;28;01mexcept\u001b[39;00m \u001b[38;5;167;01mException\u001b[39;00m \u001b[38;5;28;01mas\u001b[39;00m e:\n\u001b[1;32m     67\u001b[0m     filtered_tb \u001b[38;5;241m=\u001b[39m _process_traceback_frames(e\u001b[38;5;241m.\u001b[39m__traceback__)\n",
      "File \u001b[0;32m~/.local/lib/python3.9/site-packages/keras/engine/base_layer.py:1097\u001b[0m, in \u001b[0;36mLayer.__call__\u001b[0;34m(self, *args, **kwargs)\u001b[0m\n\u001b[1;32m   1092\u001b[0m     inputs \u001b[38;5;241m=\u001b[39m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_maybe_cast_inputs(inputs, input_list)\n\u001b[1;32m   1094\u001b[0m \u001b[38;5;28;01mwith\u001b[39;00m autocast_variable\u001b[38;5;241m.\u001b[39menable_auto_cast_variables(\n\u001b[1;32m   1095\u001b[0m     \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_compute_dtype_object\n\u001b[1;32m   1096\u001b[0m ):\n\u001b[0;32m-> 1097\u001b[0m     outputs \u001b[38;5;241m=\u001b[39m \u001b[43mcall_fn\u001b[49m\u001b[43m(\u001b[49m\u001b[43minputs\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43margs\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43mkwargs\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m   1099\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_activity_regularizer:\n\u001b[1;32m   1100\u001b[0m     \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_handle_activity_regularization(inputs, outputs)\n",
      "File \u001b[0;32m~/.local/lib/python3.9/site-packages/keras/utils/traceback_utils.py:96\u001b[0m, in \u001b[0;36minject_argument_info_in_traceback.<locals>.error_handler\u001b[0;34m(*args, **kwargs)\u001b[0m\n\u001b[1;32m     94\u001b[0m bound_signature \u001b[38;5;241m=\u001b[39m \u001b[38;5;28;01mNone\u001b[39;00m\n\u001b[1;32m     95\u001b[0m \u001b[38;5;28;01mtry\u001b[39;00m:\n\u001b[0;32m---> 96\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[43mfn\u001b[49m\u001b[43m(\u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43margs\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43mkwargs\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m     97\u001b[0m \u001b[38;5;28;01mexcept\u001b[39;00m \u001b[38;5;167;01mException\u001b[39;00m \u001b[38;5;28;01mas\u001b[39;00m e:\n\u001b[1;32m     98\u001b[0m     \u001b[38;5;28;01mif\u001b[39;00m \u001b[38;5;28mhasattr\u001b[39m(e, \u001b[38;5;124m\"\u001b[39m\u001b[38;5;124m_keras_call_info_injected\u001b[39m\u001b[38;5;124m\"\u001b[39m):\n\u001b[1;32m     99\u001b[0m         \u001b[38;5;66;03m# Only inject info for the innermost failing call\u001b[39;00m\n",
      "File \u001b[0;32m~/.local/lib/python3.9/site-packages/keras/engine/functional.py:510\u001b[0m, in \u001b[0;36mFunctional.call\u001b[0;34m(self, inputs, training, mask)\u001b[0m\n\u001b[1;32m    491\u001b[0m \u001b[38;5;129m@doc_controls\u001b[39m\u001b[38;5;241m.\u001b[39mdo_not_doc_inheritable\n\u001b[1;32m    492\u001b[0m \u001b[38;5;28;01mdef\u001b[39;00m \u001b[38;5;21mcall\u001b[39m(\u001b[38;5;28mself\u001b[39m, inputs, training\u001b[38;5;241m=\u001b[39m\u001b[38;5;28;01mNone\u001b[39;00m, mask\u001b[38;5;241m=\u001b[39m\u001b[38;5;28;01mNone\u001b[39;00m):\n\u001b[1;32m    493\u001b[0m     \u001b[38;5;124;03m\"\"\"Calls the model on new inputs.\u001b[39;00m\n\u001b[1;32m    494\u001b[0m \n\u001b[1;32m    495\u001b[0m \u001b[38;5;124;03m    In this case `call` just reapplies\u001b[39;00m\n\u001b[0;32m   (...)\u001b[0m\n\u001b[1;32m    508\u001b[0m \u001b[38;5;124;03m        a list of tensors if there are more than one outputs.\u001b[39;00m\n\u001b[1;32m    509\u001b[0m \u001b[38;5;124;03m    \"\"\"\u001b[39;00m\n\u001b[0;32m--> 510\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43m_run_internal_graph\u001b[49m\u001b[43m(\u001b[49m\u001b[43minputs\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mtraining\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mtraining\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mmask\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mmask\u001b[49m\u001b[43m)\u001b[49m\n",
      "File \u001b[0;32m~/.local/lib/python3.9/site-packages/keras/engine/functional.py:667\u001b[0m, in \u001b[0;36mFunctional._run_internal_graph\u001b[0;34m(self, inputs, training, mask)\u001b[0m\n\u001b[1;32m    664\u001b[0m     \u001b[38;5;28;01mcontinue\u001b[39;00m  \u001b[38;5;66;03m# Node is not computable, try skipping.\u001b[39;00m\n\u001b[1;32m    666\u001b[0m args, kwargs \u001b[38;5;241m=\u001b[39m node\u001b[38;5;241m.\u001b[39mmap_arguments(tensor_dict)\n\u001b[0;32m--> 667\u001b[0m outputs \u001b[38;5;241m=\u001b[39m \u001b[43mnode\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mlayer\u001b[49m\u001b[43m(\u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43margs\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43mkwargs\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m    669\u001b[0m \u001b[38;5;66;03m# Update tensor_dict.\u001b[39;00m\n\u001b[1;32m    670\u001b[0m \u001b[38;5;28;01mfor\u001b[39;00m x_id, y \u001b[38;5;129;01min\u001b[39;00m \u001b[38;5;28mzip\u001b[39m(\n\u001b[1;32m    671\u001b[0m     node\u001b[38;5;241m.\u001b[39mflat_output_ids, tf\u001b[38;5;241m.\u001b[39mnest\u001b[38;5;241m.\u001b[39mflatten(outputs)\n\u001b[1;32m    672\u001b[0m ):\n",
      "File \u001b[0;32m~/.local/lib/python3.9/site-packages/keras/utils/traceback_utils.py:65\u001b[0m, in \u001b[0;36mfilter_traceback.<locals>.error_handler\u001b[0;34m(*args, **kwargs)\u001b[0m\n\u001b[1;32m     63\u001b[0m filtered_tb \u001b[38;5;241m=\u001b[39m \u001b[38;5;28;01mNone\u001b[39;00m\n\u001b[1;32m     64\u001b[0m \u001b[38;5;28;01mtry\u001b[39;00m:\n\u001b[0;32m---> 65\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[43mfn\u001b[49m\u001b[43m(\u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43margs\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43mkwargs\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m     66\u001b[0m \u001b[38;5;28;01mexcept\u001b[39;00m \u001b[38;5;167;01mException\u001b[39;00m \u001b[38;5;28;01mas\u001b[39;00m e:\n\u001b[1;32m     67\u001b[0m     filtered_tb \u001b[38;5;241m=\u001b[39m _process_traceback_frames(e\u001b[38;5;241m.\u001b[39m__traceback__)\n",
      "File \u001b[0;32m~/.local/lib/python3.9/site-packages/keras/engine/base_layer.py:1097\u001b[0m, in \u001b[0;36mLayer.__call__\u001b[0;34m(self, *args, **kwargs)\u001b[0m\n\u001b[1;32m   1092\u001b[0m     inputs \u001b[38;5;241m=\u001b[39m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_maybe_cast_inputs(inputs, input_list)\n\u001b[1;32m   1094\u001b[0m \u001b[38;5;28;01mwith\u001b[39;00m autocast_variable\u001b[38;5;241m.\u001b[39menable_auto_cast_variables(\n\u001b[1;32m   1095\u001b[0m     \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_compute_dtype_object\n\u001b[1;32m   1096\u001b[0m ):\n\u001b[0;32m-> 1097\u001b[0m     outputs \u001b[38;5;241m=\u001b[39m \u001b[43mcall_fn\u001b[49m\u001b[43m(\u001b[49m\u001b[43minputs\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43margs\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43mkwargs\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m   1099\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_activity_regularizer:\n\u001b[1;32m   1100\u001b[0m     \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_handle_activity_regularization(inputs, outputs)\n",
      "File \u001b[0;32m~/.local/lib/python3.9/site-packages/keras/utils/traceback_utils.py:96\u001b[0m, in \u001b[0;36minject_argument_info_in_traceback.<locals>.error_handler\u001b[0;34m(*args, **kwargs)\u001b[0m\n\u001b[1;32m     94\u001b[0m bound_signature \u001b[38;5;241m=\u001b[39m \u001b[38;5;28;01mNone\u001b[39;00m\n\u001b[1;32m     95\u001b[0m \u001b[38;5;28;01mtry\u001b[39;00m:\n\u001b[0;32m---> 96\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[43mfn\u001b[49m\u001b[43m(\u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43margs\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43mkwargs\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m     97\u001b[0m \u001b[38;5;28;01mexcept\u001b[39;00m \u001b[38;5;167;01mException\u001b[39;00m \u001b[38;5;28;01mas\u001b[39;00m e:\n\u001b[1;32m     98\u001b[0m     \u001b[38;5;28;01mif\u001b[39;00m \u001b[38;5;28mhasattr\u001b[39m(e, \u001b[38;5;124m\"\u001b[39m\u001b[38;5;124m_keras_call_info_injected\u001b[39m\u001b[38;5;124m\"\u001b[39m):\n\u001b[1;32m     99\u001b[0m         \u001b[38;5;66;03m# Only inject info for the innermost failing call\u001b[39;00m\n",
      "File \u001b[0;32m~/.local/lib/python3.9/site-packages/keras/layers/convolutional/conv3d_transpose.py:314\u001b[0m, in \u001b[0;36mConv3DTranspose.call\u001b[0;34m(self, inputs)\u001b[0m\n\u001b[1;32m    311\u001b[0m     strides \u001b[38;5;241m=\u001b[39m (\u001b[38;5;241m1\u001b[39m, stride_d, stride_h, stride_w, \u001b[38;5;241m1\u001b[39m)\n\u001b[1;32m    313\u001b[0m output_shape_tensor \u001b[38;5;241m=\u001b[39m tf\u001b[38;5;241m.\u001b[39mstack(output_shape)\n\u001b[0;32m--> 314\u001b[0m outputs \u001b[38;5;241m=\u001b[39m \u001b[43mtf\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mnn\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mconv3d_transpose\u001b[49m\u001b[43m(\u001b[49m\n\u001b[1;32m    315\u001b[0m \u001b[43m    \u001b[49m\u001b[43minputs\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m    316\u001b[0m \u001b[43m    \u001b[49m\u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mkernel\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m    317\u001b[0m \u001b[43m    \u001b[49m\u001b[43moutput_shape_tensor\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m    318\u001b[0m \u001b[43m    \u001b[49m\u001b[43mstrides\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m    319\u001b[0m \u001b[43m    \u001b[49m\u001b[43mdata_format\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mconv_utils\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mconvert_data_format\u001b[49m\u001b[43m(\u001b[49m\n\u001b[1;32m    320\u001b[0m \u001b[43m        \u001b[49m\u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mdata_format\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mndim\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[38;5;241;43m5\u001b[39;49m\n\u001b[1;32m    321\u001b[0m \u001b[43m    \u001b[49m\u001b[43m)\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m    322\u001b[0m \u001b[43m    \u001b[49m\u001b[43mpadding\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mpadding\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mupper\u001b[49m\u001b[43m(\u001b[49m\u001b[43m)\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m    323\u001b[0m \u001b[43m\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m    325\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m \u001b[38;5;129;01mnot\u001b[39;00m tf\u001b[38;5;241m.\u001b[39mexecuting_eagerly():\n\u001b[1;32m    326\u001b[0m     \u001b[38;5;66;03m# Infer the static output shape:\u001b[39;00m\n\u001b[1;32m    327\u001b[0m     out_shape \u001b[38;5;241m=\u001b[39m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mcompute_output_shape(inputs\u001b[38;5;241m.\u001b[39mshape)\n",
      "File \u001b[0;32m~/.local/lib/python3.9/site-packages/tensorflow/python/util/traceback_utils.py:150\u001b[0m, in \u001b[0;36mfilter_traceback.<locals>.error_handler\u001b[0;34m(*args, **kwargs)\u001b[0m\n\u001b[1;32m    148\u001b[0m filtered_tb \u001b[38;5;241m=\u001b[39m \u001b[38;5;28;01mNone\u001b[39;00m\n\u001b[1;32m    149\u001b[0m \u001b[38;5;28;01mtry\u001b[39;00m:\n\u001b[0;32m--> 150\u001b[0m   \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[43mfn\u001b[49m\u001b[43m(\u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43margs\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43mkwargs\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m    151\u001b[0m \u001b[38;5;28;01mexcept\u001b[39;00m \u001b[38;5;167;01mException\u001b[39;00m \u001b[38;5;28;01mas\u001b[39;00m e:\n\u001b[1;32m    152\u001b[0m   filtered_tb \u001b[38;5;241m=\u001b[39m _process_traceback_frames(e\u001b[38;5;241m.\u001b[39m__traceback__)\n",
      "File \u001b[0;32m~/.local/lib/python3.9/site-packages/tensorflow/python/util/dispatch.py:1176\u001b[0m, in \u001b[0;36madd_dispatch_support.<locals>.decorator.<locals>.op_dispatch_handler\u001b[0;34m(*args, **kwargs)\u001b[0m\n\u001b[1;32m   1174\u001b[0m \u001b[38;5;66;03m# Fallback dispatch system (dispatch v1):\u001b[39;00m\n\u001b[1;32m   1175\u001b[0m \u001b[38;5;28;01mtry\u001b[39;00m:\n\u001b[0;32m-> 1176\u001b[0m   \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[43mdispatch_target\u001b[49m\u001b[43m(\u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43margs\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43mkwargs\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m   1177\u001b[0m \u001b[38;5;28;01mexcept\u001b[39;00m (\u001b[38;5;167;01mTypeError\u001b[39;00m, \u001b[38;5;167;01mValueError\u001b[39;00m):\n\u001b[1;32m   1178\u001b[0m   \u001b[38;5;66;03m# Note: convert_to_eager_tensor currently raises a ValueError, not a\u001b[39;00m\n\u001b[1;32m   1179\u001b[0m   \u001b[38;5;66;03m# TypeError, when given unexpected types.  So we need to catch both.\u001b[39;00m\n\u001b[1;32m   1180\u001b[0m   result \u001b[38;5;241m=\u001b[39m dispatch(op_dispatch_handler, args, kwargs)\n",
      "File \u001b[0;32m~/.local/lib/python3.9/site-packages/tensorflow/python/ops/nn_ops.py:3404\u001b[0m, in \u001b[0;36mconv3d_transpose_v2\u001b[0;34m(input, filters, output_shape, strides, padding, data_format, dilations, name)\u001b[0m\n\u001b[1;32m   3401\u001b[0m strides \u001b[38;5;241m=\u001b[39m _get_sequence(strides, \u001b[38;5;241m3\u001b[39m, channel_index, \u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mstrides\u001b[39m\u001b[38;5;124m\"\u001b[39m)\n\u001b[1;32m   3402\u001b[0m dilations \u001b[38;5;241m=\u001b[39m _get_sequence(dilations, \u001b[38;5;241m3\u001b[39m, channel_index, \u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mdilations\u001b[39m\u001b[38;5;124m\"\u001b[39m)\n\u001b[0;32m-> 3404\u001b[0m \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[43mgen_nn_ops\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mconv3d_backprop_input_v2\u001b[49m\u001b[43m(\u001b[49m\n\u001b[1;32m   3405\u001b[0m \u001b[43m    \u001b[49m\u001b[43minput_sizes\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43moutput_shape\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m   3406\u001b[0m \u001b[43m    \u001b[49m\u001b[38;5;28;43mfilter\u001b[39;49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mfilters\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m   3407\u001b[0m \u001b[43m    \u001b[49m\u001b[43mout_backprop\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[38;5;28;43minput\u001b[39;49m\u001b[43m,\u001b[49m\n\u001b[1;32m   3408\u001b[0m \u001b[43m    \u001b[49m\u001b[43mstrides\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mstrides\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m   3409\u001b[0m \u001b[43m    \u001b[49m\u001b[43mpadding\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mpadding\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m   3410\u001b[0m \u001b[43m    \u001b[49m\u001b[43mdata_format\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mdata_format\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m   3411\u001b[0m \u001b[43m    \u001b[49m\u001b[43mdilations\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mdilations\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m   3412\u001b[0m \u001b[43m    \u001b[49m\u001b[43mname\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mname\u001b[49m\u001b[43m)\u001b[49m\n",
      "File \u001b[0;32m~/.local/lib/python3.9/site-packages/tensorflow/python/ops/gen_nn_ops.py:2075\u001b[0m, in \u001b[0;36mconv3d_backprop_input_v2\u001b[0;34m(input_sizes, filter, out_backprop, strides, padding, data_format, dilations, name)\u001b[0m\n\u001b[1;32m   2073\u001b[0m   \u001b[38;5;28;01mpass\u001b[39;00m\n\u001b[1;32m   2074\u001b[0m \u001b[38;5;28;01mtry\u001b[39;00m:\n\u001b[0;32m-> 2075\u001b[0m   \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[43mconv3d_backprop_input_v2_eager_fallback\u001b[49m\u001b[43m(\u001b[49m\n\u001b[1;32m   2076\u001b[0m \u001b[43m      \u001b[49m\u001b[43minput_sizes\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;28;43mfilter\u001b[39;49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mout_backprop\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mstrides\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mstrides\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mpadding\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mpadding\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m   2077\u001b[0m \u001b[43m      \u001b[49m\u001b[43mdata_format\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mdata_format\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mdilations\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mdilations\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mname\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mname\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mctx\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43m_ctx\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m   2078\u001b[0m \u001b[38;5;28;01mexcept\u001b[39;00m _core\u001b[38;5;241m.\u001b[39m_SymbolicException:\n\u001b[1;32m   2079\u001b[0m   \u001b[38;5;28;01mpass\u001b[39;00m  \u001b[38;5;66;03m# Add nodes to the TensorFlow graph.\u001b[39;00m\n",
      "File \u001b[0;32m~/.local/lib/python3.9/site-packages/tensorflow/python/ops/gen_nn_ops.py:2141\u001b[0m, in \u001b[0;36mconv3d_backprop_input_v2_eager_fallback\u001b[0;34m(input_sizes, filter, out_backprop, strides, padding, data_format, dilations, name, ctx)\u001b[0m\n\u001b[1;32m   2138\u001b[0m _inputs_flat \u001b[38;5;241m=\u001b[39m [input_sizes, \u001b[38;5;28mfilter\u001b[39m, out_backprop]\n\u001b[1;32m   2139\u001b[0m _attrs \u001b[38;5;241m=\u001b[39m (\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mT\u001b[39m\u001b[38;5;124m\"\u001b[39m, _attr_T, \u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mstrides\u001b[39m\u001b[38;5;124m\"\u001b[39m, strides, \u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mpadding\u001b[39m\u001b[38;5;124m\"\u001b[39m, padding,\n\u001b[1;32m   2140\u001b[0m \u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mdata_format\u001b[39m\u001b[38;5;124m\"\u001b[39m, data_format, \u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mdilations\u001b[39m\u001b[38;5;124m\"\u001b[39m, dilations, \u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mTshape\u001b[39m\u001b[38;5;124m\"\u001b[39m, _attr_Tshape)\n\u001b[0;32m-> 2141\u001b[0m _result \u001b[38;5;241m=\u001b[39m \u001b[43m_execute\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mexecute\u001b[49m\u001b[43m(\u001b[49m\u001b[38;5;124;43mb\u001b[39;49m\u001b[38;5;124;43m\"\u001b[39;49m\u001b[38;5;124;43mConv3DBackpropInputV2\u001b[39;49m\u001b[38;5;124;43m\"\u001b[39;49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m1\u001b[39;49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43minputs\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43m_inputs_flat\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m   2142\u001b[0m \u001b[43m                           \u001b[49m\u001b[43mattrs\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43m_attrs\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mctx\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mctx\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mname\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mname\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m   2143\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m _execute\u001b[38;5;241m.\u001b[39mmust_record_gradient():\n\u001b[1;32m   2144\u001b[0m   _execute\u001b[38;5;241m.\u001b[39mrecord_gradient(\n\u001b[1;32m   2145\u001b[0m       \u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mConv3DBackpropInputV2\u001b[39m\u001b[38;5;124m\"\u001b[39m, _inputs_flat, _attrs, _result)\n",
      "File \u001b[0;32m~/.local/lib/python3.9/site-packages/tensorflow/python/eager/execute.py:54\u001b[0m, in \u001b[0;36mquick_execute\u001b[0;34m(op_name, num_outputs, inputs, attrs, ctx, name)\u001b[0m\n\u001b[1;32m     52\u001b[0m \u001b[38;5;28;01mtry\u001b[39;00m:\n\u001b[1;32m     53\u001b[0m   ctx\u001b[38;5;241m.\u001b[39mensure_initialized()\n\u001b[0;32m---> 54\u001b[0m   tensors \u001b[38;5;241m=\u001b[39m \u001b[43mpywrap_tfe\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mTFE_Py_Execute\u001b[49m\u001b[43m(\u001b[49m\u001b[43mctx\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43m_handle\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mdevice_name\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mop_name\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m     55\u001b[0m \u001b[43m                                      \u001b[49m\u001b[43minputs\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mattrs\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mnum_outputs\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m     56\u001b[0m \u001b[38;5;28;01mexcept\u001b[39;00m core\u001b[38;5;241m.\u001b[39m_NotOkStatusException \u001b[38;5;28;01mas\u001b[39;00m e:\n\u001b[1;32m     57\u001b[0m   \u001b[38;5;28;01mif\u001b[39;00m name \u001b[38;5;129;01mis\u001b[39;00m \u001b[38;5;129;01mnot\u001b[39;00m \u001b[38;5;28;01mNone\u001b[39;00m:\n",
      "\u001b[0;31mKeyboardInterrupt\u001b[0m: "
     ]
    }
   ],
   "source": [
    "batch_size = 1\n",
    "for i in range(2, 10):\n",
    "    seq_len    = 32 * i\n",
    "    print(seq_len)\n",
    "    show_memory('Before')\n",
    "\n",
    "    inp = tf.random.uniform((batch_size, seq_len) + tuple(model.input_shape[2:]), 0., 1.)\n",
    "    _ = model(inp)\n",
    "\n",
    "    show_memory('After fit')\n",
    "    tf.config.experimental.reset_memory_stats('GPU:0')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "b8a264d2-42dd-44d8-9c22-bdfd50bd662e",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "76.8\n"
     ]
    }
   ],
   "source": [
    "print(256 * 3 / 10)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f275c9ef-a0d3-426c-9418-c9bf5d6c66d1",
   "metadata": {},
   "outputs": [],
   "source": [
    "model.summary()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9d08fefe-0349-44a1-bf71-decaff7fad79",
   "metadata": {
    "jp-MarkdownHeadingCollapsed": true,
    "tags": []
   },
   "source": [
    "## Test U-Net 3D"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "d3261100-9a3a-47f3-8f8a-16c47af78fc2",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "2023-03-21 10:47:32.945584: I tensorflow/core/platform/cpu_feature_guard.cc:193] This TensorFlow binary is optimized with oneAPI Deep Neural Network Library (oneDNN) to use the following CPU instructions in performance-critical operations:  AVX2 AVX512F AVX512_VNNI FMA\n",
      "To enable them in other operations, rebuild TensorFlow with the appropriate compiler flags.\n",
      "2023-03-21 10:47:33.323590: I tensorflow/core/common_runtime/gpu/gpu_device.cc:1616] Created device /job:localhost/replica:0/task:0/device:GPU:0 with 14783 MB memory:  -> device: 0, name: Quadro RTX 5000, pci bus id: 0000:17:00.0, compute capability: 7.5\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "5823497\n",
      "{'current': '0.022', 'peak': '0.026'}\n"
     ]
    }
   ],
   "source": [
    "importlib.reload(unet_arch)\n",
    "\n",
    "model = unet_arch.UNet(\n",
    "    input_shape = (128, 512, 512, 1), output_dim  = 1, n_stages    = 6, n_conv_per_stage = [1, 1, 1, 2, 2, 2],\n",
    "    drop_rate = lambda i: 0. if i == 5 else 0.25,\n",
    "    concat_mode = 'add',\n",
    "    pool_strides = [2, 2, 2, 2, 2],\n",
    "    filters = [8, 16, 32, 64, 128, 256]\n",
    ")\n",
    "print(model.count_params())\n",
    "model.compile(loss = 'binary_crossentropy', optimizer = 'adam')\n",
    "\n",
    "show_memory()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "8c48f2f9-4940-48c6-ab8f-6c72cef5c812",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Before\t: {'current': '0.022', 'peak': '0.026'}\n",
      "Epoch 1/5\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "2023-03-21 10:47:35.758813: I tensorflow/stream_executor/cuda/cuda_dnn.cc:384] Loaded cuDNN version 8204\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "2/2 [==============================] - 32s 2s/step - loss: 0.7371\n",
      "Epoch 2/5\n",
      "2/2 [==============================] - 4s 2s/step - loss: 0.7000\n",
      "Epoch 3/5\n",
      "2/2 [==============================] - 4s 2s/step - loss: 0.6798\n",
      "Epoch 4/5\n",
      "2/2 [==============================] - 4s 2s/step - loss: 0.6471\n",
      "Epoch 5/5\n",
      "2/2 [==============================] - 4s 2s/step - loss: 0.5255\n",
      "After fit\t: {'current': '0.905', 'peak': '13.136'}\n"
     ]
    }
   ],
   "source": [
    "batch_size = 1\n",
    "\n",
    "show_memory('Before')\n",
    "\n",
    "inp = tf.random.uniform((batch_size * 2, ) + tuple(model.input_shape[1:]), 0., 1.)\n",
    "out = tf.ones_like(inp)\n",
    "\n",
    "_ = model.fit(inp, out, epochs = 5, batch_size = batch_size)\n",
    "\n",
    "show_memory('After fit')\n",
    "tf.config.experimental.reset_memory_stats('GPU:0')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "77b183f7-c75c-4416-94bf-b714dac23fe3",
   "metadata": {
    "collapsed": true,
    "jupyter": {
     "outputs_hidden": true
    },
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Model: \"model\"\n",
      "__________________________________________________________________________________________________\n",
      " Layer (type)                   Output Shape         Param #     Connected to                     \n",
      "==================================================================================================\n",
      " input_image (InputLayer)       [(None, 128, 512, 5  0           []                               \n",
      "                                12, 1)]                                                           \n",
      "                                                                                                  \n",
      " down_conv1 (Conv3D)            (None, 128, 512, 51  224         ['input_image[0][0]']            \n",
      "                                2, 8)                                                             \n",
      "                                                                                                  \n",
      " activation (Activation)        (None, 128, 512, 51  0           ['down_conv1[0][0]']             \n",
      "                                2, 8)                                                             \n",
      "                                                                                                  \n",
      " max_pooling3d (MaxPooling3D)   (None, 64, 256, 256  0           ['activation[0][0]']             \n",
      "                                , 8)                                                              \n",
      "                                                                                                  \n",
      " dropout (Dropout)              (None, 64, 256, 256  0           ['max_pooling3d[0][0]']          \n",
      "                                , 8)                                                              \n",
      "                                                                                                  \n",
      " down_conv2 (Conv3D)            (None, 64, 256, 256  3472        ['dropout[0][0]']                \n",
      "                                , 16)                                                             \n",
      "                                                                                                  \n",
      " activation_1 (Activation)      (None, 64, 256, 256  0           ['down_conv2[0][0]']             \n",
      "                                , 16)                                                             \n",
      "                                                                                                  \n",
      " max_pooling3d_1 (MaxPooling3D)  (None, 32, 128, 128  0          ['activation_1[0][0]']           \n",
      "                                , 16)                                                             \n",
      "                                                                                                  \n",
      " dropout_1 (Dropout)            (None, 32, 128, 128  0           ['max_pooling3d_1[0][0]']        \n",
      "                                , 16)                                                             \n",
      "                                                                                                  \n",
      " down_conv3 (Conv3D)            (None, 32, 128, 128  13856       ['dropout_1[0][0]']              \n",
      "                                , 32)                                                             \n",
      "                                                                                                  \n",
      " activation_2 (Activation)      (None, 32, 128, 128  0           ['down_conv3[0][0]']             \n",
      "                                , 32)                                                             \n",
      "                                                                                                  \n",
      " max_pooling3d_2 (MaxPooling3D)  (None, 16, 64, 64,   0          ['activation_2[0][0]']           \n",
      "                                32)                                                               \n",
      "                                                                                                  \n",
      " dropout_2 (Dropout)            (None, 16, 64, 64,   0           ['max_pooling3d_2[0][0]']        \n",
      "                                32)                                                               \n",
      "                                                                                                  \n",
      " down_conv4_1 (Conv3D)          (None, 16, 64, 64,   55360       ['dropout_2[0][0]']              \n",
      "                                64)                                                               \n",
      "                                                                                                  \n",
      " activation_3 (Activation)      (None, 16, 64, 64,   0           ['down_conv4_1[0][0]']           \n",
      "                                64)                                                               \n",
      "                                                                                                  \n",
      " down_conv4_2 (Conv3D)          (None, 16, 64, 64,   110656      ['activation_3[0][0]']           \n",
      "                                64)                                                               \n",
      "                                                                                                  \n",
      " activation_4 (Activation)      (None, 16, 64, 64,   0           ['down_conv4_2[0][0]']           \n",
      "                                64)                                                               \n",
      "                                                                                                  \n",
      " max_pooling3d_3 (MaxPooling3D)  (None, 8, 32, 32, 6  0          ['activation_4[0][0]']           \n",
      "                                4)                                                                \n",
      "                                                                                                  \n",
      " dropout_3 (Dropout)            (None, 8, 32, 32, 6  0           ['max_pooling3d_3[0][0]']        \n",
      "                                4)                                                                \n",
      "                                                                                                  \n",
      " down_conv5_1 (Conv3D)          (None, 8, 32, 32, 1  221312      ['dropout_3[0][0]']              \n",
      "                                28)                                                               \n",
      "                                                                                                  \n",
      " activation_5 (Activation)      (None, 8, 32, 32, 1  0           ['down_conv5_1[0][0]']           \n",
      "                                28)                                                               \n",
      "                                                                                                  \n",
      " down_conv5_2 (Conv3D)          (None, 8, 32, 32, 1  442496      ['activation_5[0][0]']           \n",
      "                                28)                                                               \n",
      "                                                                                                  \n",
      " activation_6 (Activation)      (None, 8, 32, 32, 1  0           ['down_conv5_2[0][0]']           \n",
      "                                28)                                                               \n",
      "                                                                                                  \n",
      " max_pooling3d_4 (MaxPooling3D)  (None, 4, 16, 16, 1  0          ['activation_6[0][0]']           \n",
      "                                28)                                                               \n",
      "                                                                                                  \n",
      " dropout_4 (Dropout)            (None, 4, 16, 16, 1  0           ['max_pooling3d_4[0][0]']        \n",
      "                                28)                                                               \n",
      "                                                                                                  \n",
      " down_conv6_1 (Conv3D)          (None, 4, 16, 16, 2  884992      ['dropout_4[0][0]']              \n",
      "                                56)                                                               \n",
      "                                                                                                  \n",
      " activation_7 (Activation)      (None, 4, 16, 16, 2  0           ['down_conv6_1[0][0]']           \n",
      "                                56)                                                               \n",
      "                                                                                                  \n",
      " down_conv6_2 (Conv3D)          (None, 4, 16, 16, 2  1769728     ['activation_7[0][0]']           \n",
      "                                56)                                                               \n",
      "                                                                                                  \n",
      " activation_8 (Activation)      (None, 4, 16, 16, 2  0           ['down_conv6_2[0][0]']           \n",
      "                                56)                                                               \n",
      "                                                                                                  \n",
      " upsampling_5 (Conv3DTranspose)  (None, 8, 32, 32, 1  884864     ['activation_8[0][0]']           \n",
      "                                28)                                                               \n",
      "                                                                                                  \n",
      " add (Add)                      (None, 8, 32, 32, 1  0           ['upsampling_5[0][0]',           \n",
      "                                28)                               'activation_6[0][0]']           \n",
      "                                                                                                  \n",
      " up_conv5_1 (Conv3D)            (None, 8, 32, 32, 1  442496      ['add[0][0]']                    \n",
      "                                28)                                                               \n",
      "                                                                                                  \n",
      " activation_9 (Activation)      (None, 8, 32, 32, 1  0           ['up_conv5_1[0][0]']             \n",
      "                                28)                                                               \n",
      "                                                                                                  \n",
      " up_conv5_2 (Conv3D)            (None, 8, 32, 32, 1  442496      ['activation_9[0][0]']           \n",
      "                                28)                                                               \n",
      "                                                                                                  \n",
      " activation_10 (Activation)     (None, 8, 32, 32, 1  0           ['up_conv5_2[0][0]']             \n",
      "                                28)                                                               \n",
      "                                                                                                  \n",
      " dropout_5 (Dropout)            (None, 8, 32, 32, 1  0           ['activation_10[0][0]']          \n",
      "                                28)                                                               \n",
      "                                                                                                  \n",
      " upsampling_4 (Conv3DTranspose)  (None, 16, 64, 64,   221248     ['dropout_5[0][0]']              \n",
      "                                64)                                                               \n",
      "                                                                                                  \n",
      " add_1 (Add)                    (None, 16, 64, 64,   0           ['upsampling_4[0][0]',           \n",
      "                                64)                               'activation_4[0][0]']           \n",
      "                                                                                                  \n",
      " up_conv4_1 (Conv3D)            (None, 16, 64, 64,   110656      ['add_1[0][0]']                  \n",
      "                                64)                                                               \n",
      "                                                                                                  \n",
      " activation_11 (Activation)     (None, 16, 64, 64,   0           ['up_conv4_1[0][0]']             \n",
      "                                64)                                                               \n",
      "                                                                                                  \n",
      " up_conv4_2 (Conv3D)            (None, 16, 64, 64,   110656      ['activation_11[0][0]']          \n",
      "                                64)                                                               \n",
      "                                                                                                  \n",
      " activation_12 (Activation)     (None, 16, 64, 64,   0           ['up_conv4_2[0][0]']             \n",
      "                                64)                                                               \n",
      "                                                                                                  \n",
      " dropout_6 (Dropout)            (None, 16, 64, 64,   0           ['activation_12[0][0]']          \n",
      "                                64)                                                               \n",
      "                                                                                                  \n",
      " upsampling_3 (Conv3DTranspose)  (None, 32, 128, 128  55328      ['dropout_6[0][0]']              \n",
      "                                , 32)                                                             \n",
      "                                                                                                  \n",
      " add_2 (Add)                    (None, 32, 128, 128  0           ['upsampling_3[0][0]',           \n",
      "                                , 32)                             'activation_2[0][0]']           \n",
      "                                                                                                  \n",
      " up_conv3 (Conv3D)              (None, 32, 128, 128  27680       ['add_2[0][0]']                  \n",
      "                                , 32)                                                             \n",
      "                                                                                                  \n",
      " activation_13 (Activation)     (None, 32, 128, 128  0           ['up_conv3[0][0]']               \n",
      "                                , 32)                                                             \n",
      "                                                                                                  \n",
      " dropout_7 (Dropout)            (None, 32, 128, 128  0           ['activation_13[0][0]']          \n",
      "                                , 32)                                                             \n",
      "                                                                                                  \n",
      " upsampling_2 (Conv3DTranspose)  (None, 64, 256, 256  13840      ['dropout_7[0][0]']              \n",
      "                                , 16)                                                             \n",
      "                                                                                                  \n",
      " add_3 (Add)                    (None, 64, 256, 256  0           ['upsampling_2[0][0]',           \n",
      "                                , 16)                             'activation_1[0][0]']           \n",
      "                                                                                                  \n",
      " up_conv2 (Conv3D)              (None, 64, 256, 256  6928        ['add_3[0][0]']                  \n",
      "                                , 16)                                                             \n",
      "                                                                                                  \n",
      " activation_14 (Activation)     (None, 64, 256, 256  0           ['up_conv2[0][0]']               \n",
      "                                , 16)                                                             \n",
      "                                                                                                  \n",
      " dropout_8 (Dropout)            (None, 64, 256, 256  0           ['activation_14[0][0]']          \n",
      "                                , 16)                                                             \n",
      "                                                                                                  \n",
      " upsampling_1 (Conv3DTranspose)  (None, 128, 512, 51  3464       ['dropout_8[0][0]']              \n",
      "                                2, 8)                                                             \n",
      "                                                                                                  \n",
      " add_4 (Add)                    (None, 128, 512, 51  0           ['upsampling_1[0][0]',           \n",
      "                                2, 8)                             'activation[0][0]']             \n",
      "                                                                                                  \n",
      " up_conv1 (Conv3D)              (None, 128, 512, 51  1736        ['add_4[0][0]']                  \n",
      "                                2, 8)                                                             \n",
      "                                                                                                  \n",
      " activation_15 (Activation)     (None, 128, 512, 51  0           ['up_conv1[0][0]']               \n",
      "                                2, 8)                                                             \n",
      "                                                                                                  \n",
      " dropout_9 (Dropout)            (None, 128, 512, 51  0           ['activation_15[0][0]']          \n",
      "                                2, 8)                                                             \n",
      "                                                                                                  \n",
      " segmentation_layer (Conv2D)    (None, 128, 512, 51  9           ['dropout_9[0][0]']              \n",
      "                                2, 1)                                                             \n",
      "                                                                                                  \n",
      "==================================================================================================\n",
      "Total params: 5,823,497\n",
      "Trainable params: 5,823,497\n",
      "Non-trainable params: 0\n",
      "__________________________________________________________________________________________________\n"
     ]
    }
   ],
   "source": [
    "model.summary()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8c383f3c-85b4-494d-8d31-428e421ab08b",
   "metadata": {
    "jp-MarkdownHeadingCollapsed": true,
    "tags": []
   },
   "source": [
    "## Test UNet 3D with strides"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "3f7721bf-0645-4a33-bdcd-e156a0731158",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "2023-03-21 11:18:44.599095: I tensorflow/core/platform/cpu_feature_guard.cc:193] This TensorFlow binary is optimized with oneAPI Deep Neural Network Library (oneDNN) to use the following CPU instructions in performance-critical operations:  AVX2 AVX512F AVX512_VNNI FMA\n",
      "To enable them in other operations, rebuild TensorFlow with the appropriate compiler flags.\n",
      "2023-03-21 11:18:44.979307: I tensorflow/core/common_runtime/gpu/gpu_device.cc:1616] Created device /job:localhost/replica:0/task:0/device:GPU:0 with 14783 MB memory:  -> device: 0, name: Quadro RTX 5000, pci bus id: 0000:17:00.0, compute capability: 7.5\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "0 (None, 128, 512, 512, 8)\n",
      "1 (None, 64, 256, 256, 16)\n",
      "2 (None, 32, 128, 128, 32)\n",
      "3 (None, 16, 64, 64, 64)\n",
      "4 (None, 8, 32, 32, 128)\n",
      "5 (None, 4, 16, 16, 256)\n",
      "4 (None, 8, 32, 32, 128) (None, 8, 32, 32, 128) 128\n",
      "3 (None, 16, 64, 64, 64) (None, 16, 64, 64, 64) 64\n",
      "2 (None, 32, 128, 128, 32) (None, 32, 128, 128, 32) 32\n",
      "1 (None, 64, 256, 256, 16) (None, 64, 256, 256, 16) 16\n",
      "0 (None, 128, 512, 512, 8) (None, 128, 512, 512, 8) 8\n",
      "5896185\n",
      "{'current': '0.022', 'peak': '0.026'}\n"
     ]
    }
   ],
   "source": [
    "importlib.reload(unet_arch)\n",
    "\n",
    "model = unet_arch.UNet(\n",
    "    input_shape = (128, 512, 512, 1), output_dim  = 1, n_stages = 6, n_conv_per_stage = 2,\n",
    "    drop_rate = lambda i: 0. if i == 5 else 0.25,\n",
    "    concat_mode = 'add',\n",
    "    pool_type = None,\n",
    "    strides = 2,\n",
    "    filters = [8, 16, 32, 64, 128, 256]\n",
    ")\n",
    "print(model.count_params())\n",
    "model.compile(loss = 'binary_crossentropy', optimizer = 'adam')\n",
    "\n",
    "show_memory()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "2c6babd0-1639-468b-9260-573a09198e16",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Before\t: {'current': '0.022', 'peak': '0.026'}\n",
      "Epoch 1/5\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "2023-03-21 11:18:53.017038: I tensorflow/stream_executor/cuda/cuda_dnn.cc:384] Loaded cuDNN version 8204\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "2/2 [==============================] - 37s 3s/step - loss: 0.6954\n",
      "Epoch 2/5\n",
      "2/2 [==============================] - 6s 3s/step - loss: 0.6638\n",
      "Epoch 3/5\n",
      "2/2 [==============================] - 6s 3s/step - loss: 0.6255\n",
      "Epoch 4/5\n",
      "2/2 [==============================] - 6s 3s/step - loss: 0.5678\n",
      "Epoch 5/5\n",
      "2/2 [==============================] - 6s 3s/step - loss: 0.4684\n",
      "After fit\t: {'current': '0.910', 'peak': '13.031'}\n"
     ]
    }
   ],
   "source": [
    "batch_size = 1\n",
    "\n",
    "show_memory('Before')\n",
    "\n",
    "inp = tf.random.uniform((batch_size * 2, ) + tuple(model.input_shape[1:]), 0., 1.)\n",
    "out = tf.ones_like(inp)\n",
    "\n",
    "_ = model.fit(inp, out, epochs = 5, batch_size = batch_size)\n",
    "\n",
    "show_memory('After fit')\n",
    "tf.config.experimental.reset_memory_stats('GPU:0')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "2a2a7c43-030d-4f53-9d9a-a80293c9360c",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Model: \"model_2\"\n",
      "__________________________________________________________________________________________________\n",
      " Layer (type)                   Output Shape         Param #     Connected to                     \n",
      "==================================================================================================\n",
      " input_image (InputLayer)       [(None, 128, 512, 5  0           []                               \n",
      "                                12, 1)]                                                           \n",
      "                                                                                                  \n",
      " down_conv1 (Conv3D)            (None, 128, 512, 51  224         ['input_image[0][0]']            \n",
      "                                2, 8)                                                             \n",
      "                                                                                                  \n",
      " activation_123 (Activation)    (None, 128, 512, 51  0           ['down_conv1[0][0]']             \n",
      "                                2, 8)                                                             \n",
      "                                                                                                  \n",
      " down_conv1_1 (Conv3D)          (None, 64, 256, 256  1736        ['activation_123[0][0]']         \n",
      "                                , 8)                                                              \n",
      "                                                                                                  \n",
      " activation_124 (Activation)    (None, 64, 256, 256  0           ['down_conv1_1[0][0]']           \n",
      "                                , 8)                                                              \n",
      "                                                                                                  \n",
      " dropout_55 (Dropout)           (None, 64, 256, 256  0           ['activation_124[0][0]']         \n",
      "                                , 8)                                                              \n",
      "                                                                                                  \n",
      " down_conv2 (Conv3D)            (None, 64, 256, 256  3472        ['dropout_55[0][0]']             \n",
      "                                , 16)                                                             \n",
      "                                                                                                  \n",
      " activation_125 (Activation)    (None, 64, 256, 256  0           ['down_conv2[0][0]']             \n",
      "                                , 16)                                                             \n",
      "                                                                                                  \n",
      " down_conv2_1 (Conv3D)          (None, 32, 128, 128  6928        ['activation_125[0][0]']         \n",
      "                                , 16)                                                             \n",
      "                                                                                                  \n",
      " activation_126 (Activation)    (None, 32, 128, 128  0           ['down_conv2_1[0][0]']           \n",
      "                                , 16)                                                             \n",
      "                                                                                                  \n",
      " dropout_56 (Dropout)           (None, 32, 128, 128  0           ['activation_126[0][0]']         \n",
      "                                , 16)                                                             \n",
      "                                                                                                  \n",
      " down_conv3 (Conv3D)            (None, 32, 128, 128  13856       ['dropout_56[0][0]']             \n",
      "                                , 32)                                                             \n",
      "                                                                                                  \n",
      " activation_127 (Activation)    (None, 32, 128, 128  0           ['down_conv3[0][0]']             \n",
      "                                , 32)                                                             \n",
      "                                                                                                  \n",
      " down_conv3_1 (Conv3D)          (None, 16, 64, 64,   27680       ['activation_127[0][0]']         \n",
      "                                32)                                                               \n",
      "                                                                                                  \n",
      " activation_128 (Activation)    (None, 16, 64, 64,   0           ['down_conv3_1[0][0]']           \n",
      "                                32)                                                               \n",
      "                                                                                                  \n",
      " dropout_57 (Dropout)           (None, 16, 64, 64,   0           ['activation_128[0][0]']         \n",
      "                                32)                                                               \n",
      "                                                                                                  \n",
      " down_conv4 (Conv3D)            (None, 16, 64, 64,   55360       ['dropout_57[0][0]']             \n",
      "                                64)                                                               \n",
      "                                                                                                  \n",
      " activation_129 (Activation)    (None, 16, 64, 64,   0           ['down_conv4[0][0]']             \n",
      "                                64)                                                               \n",
      "                                                                                                  \n",
      " down_conv4_1 (Conv3D)          (None, 8, 32, 32, 6  110656      ['activation_129[0][0]']         \n",
      "                                4)                                                                \n",
      "                                                                                                  \n",
      " activation_130 (Activation)    (None, 8, 32, 32, 6  0           ['down_conv4_1[0][0]']           \n",
      "                                4)                                                                \n",
      "                                                                                                  \n",
      " dropout_58 (Dropout)           (None, 8, 32, 32, 6  0           ['activation_130[0][0]']         \n",
      "                                4)                                                                \n",
      "                                                                                                  \n",
      " down_conv5 (Conv3D)            (None, 8, 32, 32, 1  221312      ['dropout_58[0][0]']             \n",
      "                                28)                                                               \n",
      "                                                                                                  \n",
      " activation_131 (Activation)    (None, 8, 32, 32, 1  0           ['down_conv5[0][0]']             \n",
      "                                28)                                                               \n",
      "                                                                                                  \n",
      " down_conv5_1 (Conv3D)          (None, 4, 16, 16, 1  442496      ['activation_131[0][0]']         \n",
      "                                28)                                                               \n",
      "                                                                                                  \n",
      " activation_132 (Activation)    (None, 4, 16, 16, 1  0           ['down_conv5_1[0][0]']           \n",
      "                                28)                                                               \n",
      "                                                                                                  \n",
      " dropout_59 (Dropout)           (None, 4, 16, 16, 1  0           ['activation_132[0][0]']         \n",
      "                                28)                                                               \n",
      "                                                                                                  \n",
      " down_conv6_1 (Conv3D)          (None, 4, 16, 16, 2  884992      ['dropout_59[0][0]']             \n",
      "                                56)                                                               \n",
      "                                                                                                  \n",
      " activation_133 (Activation)    (None, 4, 16, 16, 2  0           ['down_conv6_1[0][0]']           \n",
      "                                56)                                                               \n",
      "                                                                                                  \n",
      " down_conv6_2 (Conv3D)          (None, 4, 16, 16, 2  1769728     ['activation_133[0][0]']         \n",
      "                                56)                                                               \n",
      "                                                                                                  \n",
      " activation_134 (Activation)    (None, 4, 16, 16, 2  0           ['down_conv6_2[0][0]']           \n",
      "                                56)                                                               \n",
      "                                                                                                  \n",
      " upsampling_5 (Conv3DTranspose)  (None, 8, 32, 32, 1  884864     ['activation_134[0][0]']         \n",
      "                                28)                                                               \n",
      "                                                                                                  \n",
      " add_17 (Add)                   (None, 8, 32, 32, 1  0           ['upsampling_5[0][0]',           \n",
      "                                28)                               'activation_131[0][0]']         \n",
      "                                                                                                  \n",
      " up_conv5_1 (Conv3D)            (None, 8, 32, 32, 1  442496      ['add_17[0][0]']                 \n",
      "                                28)                                                               \n",
      "                                                                                                  \n",
      " activation_135 (Activation)    (None, 8, 32, 32, 1  0           ['up_conv5_1[0][0]']             \n",
      "                                28)                                                               \n",
      "                                                                                                  \n",
      " up_conv5_2 (Conv3D)            (None, 8, 32, 32, 1  442496      ['activation_135[0][0]']         \n",
      "                                28)                                                               \n",
      "                                                                                                  \n",
      " activation_136 (Activation)    (None, 8, 32, 32, 1  0           ['up_conv5_2[0][0]']             \n",
      "                                28)                                                               \n",
      "                                                                                                  \n",
      " dropout_60 (Dropout)           (None, 8, 32, 32, 1  0           ['activation_136[0][0]']         \n",
      "                                28)                                                               \n",
      "                                                                                                  \n",
      " upsampling_4 (Conv3DTranspose)  (None, 16, 64, 64,   221248     ['dropout_60[0][0]']             \n",
      "                                64)                                                               \n",
      "                                                                                                  \n",
      " add_18 (Add)                   (None, 16, 64, 64,   0           ['upsampling_4[0][0]',           \n",
      "                                64)                               'activation_129[0][0]']         \n",
      "                                                                                                  \n",
      " up_conv4_1 (Conv3D)            (None, 16, 64, 64,   110656      ['add_18[0][0]']                 \n",
      "                                64)                                                               \n",
      "                                                                                                  \n",
      " activation_137 (Activation)    (None, 16, 64, 64,   0           ['up_conv4_1[0][0]']             \n",
      "                                64)                                                               \n",
      "                                                                                                  \n",
      " up_conv4_2 (Conv3D)            (None, 16, 64, 64,   110656      ['activation_137[0][0]']         \n",
      "                                64)                                                               \n",
      "                                                                                                  \n",
      " activation_138 (Activation)    (None, 16, 64, 64,   0           ['up_conv4_2[0][0]']             \n",
      "                                64)                                                               \n",
      "                                                                                                  \n",
      " dropout_61 (Dropout)           (None, 16, 64, 64,   0           ['activation_138[0][0]']         \n",
      "                                64)                                                               \n",
      "                                                                                                  \n",
      " upsampling_3 (Conv3DTranspose)  (None, 32, 128, 128  55328      ['dropout_61[0][0]']             \n",
      "                                , 32)                                                             \n",
      "                                                                                                  \n",
      " add_19 (Add)                   (None, 32, 128, 128  0           ['upsampling_3[0][0]',           \n",
      "                                , 32)                             'activation_127[0][0]']         \n",
      "                                                                                                  \n",
      " up_conv3_1 (Conv3D)            (None, 32, 128, 128  27680       ['add_19[0][0]']                 \n",
      "                                , 32)                                                             \n",
      "                                                                                                  \n",
      " activation_139 (Activation)    (None, 32, 128, 128  0           ['up_conv3_1[0][0]']             \n",
      "                                , 32)                                                             \n",
      "                                                                                                  \n",
      " up_conv3_2 (Conv3D)            (None, 32, 128, 128  27680       ['activation_139[0][0]']         \n",
      "                                , 32)                                                             \n",
      "                                                                                                  \n",
      " activation_140 (Activation)    (None, 32, 128, 128  0           ['up_conv3_2[0][0]']             \n",
      "                                , 32)                                                             \n",
      "                                                                                                  \n",
      " dropout_62 (Dropout)           (None, 32, 128, 128  0           ['activation_140[0][0]']         \n",
      "                                , 32)                                                             \n",
      "                                                                                                  \n",
      " upsampling_2 (Conv3DTranspose)  (None, 64, 256, 256  13840      ['dropout_62[0][0]']             \n",
      "                                , 16)                                                             \n",
      "                                                                                                  \n",
      " add_20 (Add)                   (None, 64, 256, 256  0           ['upsampling_2[0][0]',           \n",
      "                                , 16)                             'activation_125[0][0]']         \n",
      "                                                                                                  \n",
      " up_conv2_1 (Conv3D)            (None, 64, 256, 256  6928        ['add_20[0][0]']                 \n",
      "                                , 16)                                                             \n",
      "                                                                                                  \n",
      " activation_141 (Activation)    (None, 64, 256, 256  0           ['up_conv2_1[0][0]']             \n",
      "                                , 16)                                                             \n",
      "                                                                                                  \n",
      " up_conv2_2 (Conv3D)            (None, 64, 256, 256  6928        ['activation_141[0][0]']         \n",
      "                                , 16)                                                             \n",
      "                                                                                                  \n",
      " activation_142 (Activation)    (None, 64, 256, 256  0           ['up_conv2_2[0][0]']             \n",
      "                                , 16)                                                             \n",
      "                                                                                                  \n",
      " dropout_63 (Dropout)           (None, 64, 256, 256  0           ['activation_142[0][0]']         \n",
      "                                , 16)                                                             \n",
      "                                                                                                  \n",
      " upsampling_1 (Conv3DTranspose)  (None, 128, 512, 51  3464       ['dropout_63[0][0]']             \n",
      "                                2, 8)                                                             \n",
      "                                                                                                  \n",
      " add_21 (Add)                   (None, 128, 512, 51  0           ['upsampling_1[0][0]',           \n",
      "                                2, 8)                             'activation_123[0][0]']         \n",
      "                                                                                                  \n",
      " up_conv1_1 (Conv3D)            (None, 128, 512, 51  1736        ['add_21[0][0]']                 \n",
      "                                2, 8)                                                             \n",
      "                                                                                                  \n",
      " activation_143 (Activation)    (None, 128, 512, 51  0           ['up_conv1_1[0][0]']             \n",
      "                                2, 8)                                                             \n",
      "                                                                                                  \n",
      " up_conv1_2 (Conv3D)            (None, 128, 512, 51  1736        ['activation_143[0][0]']         \n",
      "                                2, 8)                                                             \n",
      "                                                                                                  \n",
      " activation_144 (Activation)    (None, 128, 512, 51  0           ['up_conv1_2[0][0]']             \n",
      "                                2, 8)                                                             \n",
      "                                                                                                  \n",
      " dropout_64 (Dropout)           (None, 128, 512, 51  0           ['activation_144[0][0]']         \n",
      "                                2, 8)                                                             \n",
      "                                                                                                  \n",
      " segmentation_layer (Conv2D)    (None, 128, 512, 51  9           ['dropout_64[0][0]']             \n",
      "                                2, 1)                                                             \n",
      "                                                                                                  \n",
      "==================================================================================================\n",
      "Total params: 5,896,185\n",
      "Trainable params: 5,896,185\n",
      "Non-trainable params: 0\n",
      "__________________________________________________________________________________________________\n"
     ]
    }
   ],
   "source": [
    "model.summary()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f10bcdca-2a3c-49c6-9660-1fe5d2439856",
   "metadata": {
    "jp-MarkdownHeadingCollapsed": true,
    "tags": []
   },
   "source": [
    "## Test classifier AM-UNet 3D"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "2feb60e3-3cb2-473f-8c67-776a3d2676e3",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "INFO:tensorflow:Mixed precision compatibility check (mixed_float16): OK\n",
      "Your GPUs will likely run quickly with dtype policy mixed_float16 as they all have compute capability of at least 7.0\n",
      "Mixed precision compatibility check (mixed_float16): OK\n",
      "Your GPUs will likely run quickly with dtype policy mixed_float16 as they all have compute capability of at least 7.0\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "2023-03-30 15:01:57.966552: I tensorflow/core/platform/cpu_feature_guard.cc:193] This TensorFlow binary is optimized with oneAPI Deep Neural Network Library (oneDNN) to use the following CPU instructions in performance-critical operations:  AVX2 AVX512F AVX512_VNNI FMA\n",
      "To enable them in other operations, rebuild TensorFlow with the appropriate compiler flags.\n",
      "2023-03-30 15:01:58.353265: I tensorflow/core/common_runtime/gpu/gpu_device.cc:1616] Created device /job:localhost/replica:0/task:0/device:GPU:0 with 535 MB memory:  -> device: 0, name: Quadro RTX 5000, pci bus id: 0000:17:00.0, compute capability: 7.5\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "2519128\n",
      "{'current': '0.009', 'peak': '0.014'}\n"
     ]
    }
   ],
   "source": [
    "importlib.reload(unet_arch)\n",
    "\n",
    "model = unet_arch.UNet(\n",
    "    input_shape = (512, 512, None, 1),\n",
    "    output_dim  = 104,\n",
    "    final_activation = 'softmax',\n",
    "    \n",
    "    n_stages    = 4,\n",
    "    n_conv_per_stage = lambda i: 1,\n",
    "    up_n_conv_per_stage = lambda i: min(i, 1),\n",
    "    filters     = list(np.array([16, 32, 64, 128, 128])),\n",
    "    bnorm       = 'never',\n",
    "    activation  = 'relu',\n",
    "    drop_rate   = lambda i: 0. if i == 0 else 0.25,\n",
    "    \n",
    "    n_middle_stages = 2,\n",
    "    n_middle_conv   = 2,\n",
    "    middle_filters  = 64,\n",
    "    middle_bnorm    = 'never',\n",
    "    \n",
    "    concat_mode = lambda i: 'concat' if i > 0 else None,\n",
    "    \n",
    "    mixed_precision = True\n",
    ")\n",
    "print(model.count_params())\n",
    "model.compile(loss = 'sparse_categorical_crossentropy', optimizer = 'adam')\n",
    "\n",
    "show_memory()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "74df681a-f63d-462f-a58a-5505294ed803",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Before\t: {'current': '0.009', 'peak': '0.014'}\n",
      "Input shape : (2, 48, 512, 512, 1) - Output shape : (2, 48, 512, 512)\n",
      "Epoch 1/5\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "2023-03-29 15:38:52.065510: I tensorflow/stream_executor/cuda/cuda_dnn.cc:384] Loaded cuDNN version 8204\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "2/2 [==============================] - 5s 622ms/step - loss: 4.6445\n",
      "Epoch 2/5\n",
      "2/2 [==============================] - 1s 622ms/step - loss: 4.6444\n",
      "Epoch 3/5\n",
      "2/2 [==============================] - 1s 623ms/step - loss: 4.6444\n",
      "Epoch 4/5\n",
      "2/2 [==============================] - 1s 622ms/step - loss: 4.6444\n",
      "Epoch 5/5\n",
      "2/2 [==============================] - 1s 623ms/step - loss: 4.6444\n",
      "After fit\t: {'current': '0.216', 'peak': '8.909'}\n"
     ]
    }
   ],
   "source": [
    "batch_size = 1\n",
    "seq_len    = 48\n",
    "\n",
    "show_memory('Before')\n",
    "\n",
    "inp = tf.random.uniform((batch_size * 2, seq_len) + tuple(model.input_shape[2:]), 0., 1.)\n",
    "out = tf.random.uniform(inp.shape[:-1], 0, model.output_shape[-1], dtype = tf.int32)\n",
    "\n",
    "print('Input shape : {} - Output shape : {}'.format(inp.shape, out.shape))\n",
    "\n",
    "_ = model.fit(inp, out, epochs = 5, batch_size = batch_size)\n",
    "\n",
    "show_memory('After fit')\n",
    "tf.config.experimental.reset_memory_stats('GPU:0')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4f140612-60c7-448e-a209-6c3d145bb581",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "batch_size = 1\n",
    "for i in range(1, 3):\n",
    "    seq_len    = 16 * i\n",
    "    show_memory('Before')\n",
    "\n",
    "    inp = tf.random.uniform((batch_size, seq_len) + tuple(model.input_shape[2:]), 0., 1.)\n",
    "    print('Input shape  : {}'.format(inp.shape))\n",
    "    out = model(inp)\n",
    "    print('Output shape : {}'.format(inp.shape))\n",
    "    del out\n",
    "    \n",
    "    show_memory('After call')\n",
    "    tf.config.experimental.reset_memory_stats('GPU:0')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "054bd058-9ce2-43bd-8480-c986a61d03a9",
   "metadata": {
    "collapsed": true,
    "jupyter": {
     "outputs_hidden": true
    },
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Model: \"model\"\n",
      "__________________________________________________________________________________________________\n",
      " Layer (type)                   Output Shape         Param #     Connected to                     \n",
      "==================================================================================================\n",
      " input_image (InputLayer)       [(None, 512, 512, N  0           []                               \n",
      "                                one, 1)]                                                          \n",
      "                                                                                                  \n",
      " down_conv1 (Conv3D)            (None, 512, 512, No  448         ['input_image[0][0]']            \n",
      "                                ne, 16)                                                           \n",
      "                                                                                                  \n",
      " activation (Activation)        (None, 512, 512, No  0           ['down_conv1[0][0]']             \n",
      "                                ne, 16)                                                           \n",
      "                                                                                                  \n",
      " max_pooling3d (MaxPooling3D)   (None, 256, 256, No  0           ['activation[0][0]']             \n",
      "                                ne, 16)                                                           \n",
      "                                                                                                  \n",
      " down_conv2 (Conv3D)            (None, 256, 256, No  13856       ['max_pooling3d[0][0]']          \n",
      "                                ne, 32)                                                           \n",
      "                                                                                                  \n",
      " activation_1 (Activation)      (None, 256, 256, No  0           ['down_conv2[0][0]']             \n",
      "                                ne, 32)                                                           \n",
      "                                                                                                  \n",
      " max_pooling3d_1 (MaxPooling3D)  (None, 128, 128, No  0          ['activation_1[0][0]']           \n",
      "                                ne, 32)                                                           \n",
      "                                                                                                  \n",
      " dropout (Dropout)              (None, 128, 128, No  0           ['max_pooling3d_1[0][0]']        \n",
      "                                ne, 32)                                                           \n",
      "                                                                                                  \n",
      " down_conv3 (Conv3D)            (None, 128, 128, No  55360       ['dropout[0][0]']                \n",
      "                                ne, 64)                                                           \n",
      "                                                                                                  \n",
      " activation_2 (Activation)      (None, 128, 128, No  0           ['down_conv3[0][0]']             \n",
      "                                ne, 64)                                                           \n",
      "                                                                                                  \n",
      " max_pooling3d_2 (MaxPooling3D)  (None, 64, 64, None  0          ['activation_2[0][0]']           \n",
      "                                , 64)                                                             \n",
      "                                                                                                  \n",
      " dropout_1 (Dropout)            (None, 64, 64, None  0           ['max_pooling3d_2[0][0]']        \n",
      "                                , 64)                                                             \n",
      "                                                                                                  \n",
      " down_conv4 (Conv3D)            (None, 64, 64, None  221312      ['dropout_1[0][0]']              \n",
      "                                , 128)                                                            \n",
      "                                                                                                  \n",
      " activation_3 (Activation)      (None, 64, 64, None  0           ['down_conv4[0][0]']             \n",
      "                                , 128)                                                            \n",
      "                                                                                                  \n",
      " max_pooling3d_3 (MaxPooling3D)  (None, 32, 32, None  0          ['activation_3[0][0]']           \n",
      "                                , 128)                                                            \n",
      "                                                                                                  \n",
      " dropout_2 (Dropout)            (None, 32, 32, None  0           ['max_pooling3d_3[0][0]']        \n",
      "                                , 128)                                                            \n",
      "                                                                                                  \n",
      " middle_conv1-1 (Conv3D)        (None, 32, 32, None  221248      ['dropout_2[0][0]']              \n",
      "                                , 64)                                                             \n",
      "                                                                                                  \n",
      " activation_4 (Activation)      (None, 32, 32, None  0           ['middle_conv1-1[0][0]']         \n",
      "                                , 64)                                                             \n",
      "                                                                                                  \n",
      " dropout_3 (Dropout)            (None, 32, 32, None  0           ['activation_4[0][0]']           \n",
      "                                , 64)                                                             \n",
      "                                                                                                  \n",
      " middle_conv1-2 (Conv3D)        (None, 32, 32, None  110656      ['dropout_3[0][0]']              \n",
      "                                , 64)                                                             \n",
      "                                                                                                  \n",
      " activation_5 (Activation)      (None, 32, 32, None  0           ['middle_conv1-2[0][0]']         \n",
      "                                , 64)                                                             \n",
      "                                                                                                  \n",
      " dropout_4 (Dropout)            (None, 32, 32, None  0           ['activation_5[0][0]']           \n",
      "                                , 64)                                                             \n",
      "                                                                                                  \n",
      " middle_conv2-1 (Conv3D)        (None, 32, 32, None  110656      ['dropout_4[0][0]']              \n",
      "                                , 64)                                                             \n",
      "                                                                                                  \n",
      " activation_6 (Activation)      (None, 32, 32, None  0           ['middle_conv2-1[0][0]']         \n",
      "                                , 64)                                                             \n",
      "                                                                                                  \n",
      " dropout_5 (Dropout)            (None, 32, 32, None  0           ['activation_6[0][0]']           \n",
      "                                , 64)                                                             \n",
      "                                                                                                  \n",
      " middle_conv2-2 (Conv3D)        (None, 32, 32, None  110656      ['dropout_5[0][0]']              \n",
      "                                , 64)                                                             \n",
      "                                                                                                  \n",
      " activation_7 (Activation)      (None, 32, 32, None  0           ['middle_conv2-2[0][0]']         \n",
      "                                , 64)                                                             \n",
      "                                                                                                  \n",
      " dropout_6 (Dropout)            (None, 32, 32, None  0           ['activation_7[0][0]']           \n",
      "                                , 64)                                                             \n",
      "                                                                                                  \n",
      " upsampling_4 (Conv3DTranspose)  (None, 64, 64, None  221312     ['dropout_6[0][0]']              \n",
      "                                , 128)                                                            \n",
      "                                                                                                  \n",
      " concatenate (Concatenate)      (None, 64, 64, None  0           ['upsampling_4[0][0]',           \n",
      "                                , 256)                            'activation_3[0][0]']           \n",
      "                                                                                                  \n",
      " up_conv4 (Conv3D)              (None, 64, 64, None  884864      ['concatenate[0][0]']            \n",
      "                                , 128)                                                            \n",
      "                                                                                                  \n",
      " activation_8 (Activation)      (None, 64, 64, None  0           ['up_conv4[0][0]']               \n",
      "                                , 128)                                                            \n",
      "                                                                                                  \n",
      " dropout_7 (Dropout)            (None, 64, 64, None  0           ['activation_8[0][0]']           \n",
      "                                , 128)                                                            \n",
      "                                                                                                  \n",
      " upsampling_3 (Conv3DTranspose)  (None, 128, 128, No  221248     ['dropout_7[0][0]']              \n",
      "                                ne, 64)                                                           \n",
      "                                                                                                  \n",
      " concatenate_1 (Concatenate)    (None, 128, 128, No  0           ['upsampling_3[0][0]',           \n",
      "                                ne, 128)                          'activation_2[0][0]']           \n",
      "                                                                                                  \n",
      " up_conv3 (Conv3D)              (None, 128, 128, No  221248      ['concatenate_1[0][0]']          \n",
      "                                ne, 64)                                                           \n",
      "                                                                                                  \n",
      " activation_9 (Activation)      (None, 128, 128, No  0           ['up_conv3[0][0]']               \n",
      "                                ne, 64)                                                           \n",
      "                                                                                                  \n",
      " dropout_8 (Dropout)            (None, 128, 128, No  0           ['activation_9[0][0]']           \n",
      "                                ne, 64)                                                           \n",
      "                                                                                                  \n",
      " upsampling_2 (Conv3DTranspose)  (None, 256, 256, No  55328      ['dropout_8[0][0]']              \n",
      "                                ne, 32)                                                           \n",
      "                                                                                                  \n",
      " concatenate_2 (Concatenate)    (None, 256, 256, No  0           ['upsampling_2[0][0]',           \n",
      "                                ne, 64)                           'activation_1[0][0]']           \n",
      "                                                                                                  \n",
      " up_conv2 (Conv3D)              (None, 256, 256, No  55328       ['concatenate_2[0][0]']          \n",
      "                                ne, 32)                                                           \n",
      "                                                                                                  \n",
      " activation_10 (Activation)     (None, 256, 256, No  0           ['up_conv2[0][0]']               \n",
      "                                ne, 32)                                                           \n",
      "                                                                                                  \n",
      " dropout_9 (Dropout)            (None, 256, 256, No  0           ['activation_10[0][0]']          \n",
      "                                ne, 32)                                                           \n",
      "                                                                                                  \n",
      " upsampling_1 (Conv3DTranspose)  (None, 512, 512, No  13840      ['dropout_9[0][0]']              \n",
      "                                ne, 16)                                                           \n",
      "                                                                                                  \n",
      " segmentation_layer (Conv3D)    (None, 512, 512, No  1768        ['upsampling_1[0][0]']           \n",
      "                                ne, 104)                                                          \n",
      "                                                                                                  \n",
      " activation_11 (Activation)     (None, 512, 512, No  0           ['segmentation_layer[0][0]']     \n",
      "                                ne, 104)                                                          \n",
      "                                                                                                  \n",
      "==================================================================================================\n",
      "Total params: 2,519,128\n",
      "Trainable params: 2,519,128\n",
      "Non-trainable params: 0\n",
      "__________________________________________________________________________________________________\n"
     ]
    }
   ],
   "source": [
    "model.summary()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f303fb83-7fa5-42e7-b7aa-ba2ecceafbcb",
   "metadata": {
    "tags": []
   },
   "source": [
    "## Test TotalSegmentator"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "e611461c-8abd-4f50-ac90-89b85d65919e",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "2023-04-13 10:34:23.532705: I tensorflow/core/platform/cpu_feature_guard.cc:193] This TensorFlow binary is optimized with oneAPI Deep Neural Network Library (oneDNN) to use the following CPU instructions in performance-critical operations:  AVX2 AVX512F AVX512_VNNI FMA\n",
      "To enable them in other operations, rebuild TensorFlow with the appropriate compiler flags.\n",
      "2023-04-13 10:34:23.916325: I tensorflow/core/common_runtime/gpu/gpu_device.cc:1616] Created device /job:localhost/replica:0/task:0/device:GPU:0 with 14783 MB memory:  -> device: 0, name: Quadro RTX 5000, pci bus id: 0000:17:00.0, compute capability: 7.5\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "30479744\n",
      "{'current': '0.119', 'peak': '0.130'}\n"
     ]
    }
   ],
   "source": [
    "importlib.reload(totalsegmentator_arch)\n",
    "\n",
    "model = totalsegmentator_arch.TotalSegmentator(\n",
    "    input_shape = (None, None, None, 1),\n",
    "    output_dim  = 105,\n",
    "    pretrained  = None,\n",
    "    manual_padding = True,\n",
    "    \n",
    "    #n_conv_per_stage = lambda i: 1 if i <= 2 else 2,\n",
    "    #up_n_conv_per_stage = 1,\n",
    "    #filters = [16, 32, 64, 128, 128, 128],\n",
    "    drop_rate   = lambda i: 0. if i == 0 else 0.25,\n",
    "\n",
    ")\n",
    "print(model.count_params())\n",
    "model.compile(loss = 'sparse_categorical_crossentropy', optimizer = 'adam')\n",
    "\n",
    "show_memory()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "ea298436-de4e-4063-b3c0-06000a3c77d3",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Before\t: {'current': '0.402', 'peak': '0.402'}\n",
      "Input shape : (2, 48, 256, 256, 1) - Output shape : (2, 48, 256, 256)\n",
      "Epoch 1/5\n",
      "2/2 [==============================] - 18s 1s/step - loss: 4.7110\n",
      "Epoch 2/5\n",
      "2/2 [==============================] - 3s 1s/step - loss: 4.6852\n",
      "Epoch 3/5\n",
      "2/2 [==============================] - 3s 1s/step - loss: 4.6715\n",
      "Epoch 4/5\n",
      "2/2 [==============================] - 3s 1s/step - loss: 4.6652\n",
      "Epoch 5/5\n",
      "2/2 [==============================] - 3s 1s/step - loss: 4.6627\n",
      "After fit\t: {'current': '0.449', 'peak': '12.968'}\n"
     ]
    }
   ],
   "source": [
    "batch_size = 1\n",
    "seq_len    = 48\n",
    "img_size   = (256, 256, 1)\n",
    "\n",
    "show_memory('Before')\n",
    "\n",
    "inp = tf.random.uniform((batch_size * 2, seq_len) + img_size, 0., 1.)\n",
    "out = tf.random.uniform(inp.shape[:-1], 0, model.output_shape[-1], dtype = tf.int32)\n",
    "\n",
    "print('Input shape : {} - Output shape : {}'.format(inp.shape, out.shape))\n",
    "\n",
    "_ = model.fit(inp, out, epochs = 5, batch_size = batch_size)\n",
    "\n",
    "show_memory('After fit')\n",
    "tf.config.experimental.reset_memory_stats('GPU:0')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "4f5d88b0-ca61-46e3-b263-056a437cd729",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Before\t: {'current': '0.212', 'peak': '0.212'}\n",
      "Input shape  : (1, 16, 512, 512, 1)\n",
      "Output shape : (1, 16, 512, 512, 105)\n",
      "After call\t: {'current': '0.181', 'peak': '7.992'}\n",
      "Before\t: {'current': '0.181', 'peak': '0.181'}\n",
      "Input shape  : (1, 32, 512, 512, 1)\n",
      "Output shape : (1, 32, 512, 512, 105)\n",
      "After call\t: {'current': '0.196', 'peak': '8.671'}\n",
      "Before\t: {'current': '0.196', 'peak': '0.196'}\n",
      "Input shape  : (1, 48, 512, 512, 1)\n",
      "Output shape : (1, 48, 512, 512, 105)\n",
      "After call\t: {'current': '0.212', 'peak': '12.455'}\n"
     ]
    }
   ],
   "source": [
    "batch_size = 1\n",
    "img_size   = (512, 512, 1)\n",
    "for i in range(1, 4):\n",
    "    seq_len    = 16 * i\n",
    "    show_memory('Before')\n",
    "\n",
    "    inp = tf.random.uniform((batch_size, seq_len) + img_size, 0., 1.)\n",
    "    print('Input shape  : {}'.format(inp.shape))\n",
    "    out = model(inp)\n",
    "    print('Output shape : {}'.format(out.shape))\n",
    "    del out\n",
    "    \n",
    "    show_memory('After call')\n",
    "    tf.config.experimental.reset_memory_stats('GPU:0')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "1eead026-d11b-4b3b-9c51-c088a63fe7f2",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Model: \"model\"\n",
      "__________________________________________________________________________________________________\n",
      " Layer (type)                   Output Shape         Param #     Connected to                     \n",
      "==================================================================================================\n",
      " input_image (InputLayer)       [(None, None, None,  0           []                               \n",
      "                                 None, 1)]                                                        \n",
      "                                                                                                  \n",
      " zero_padding3d (ZeroPadding3D)  (None, None, None,   0          ['input_image[0][0]']            \n",
      "                                None, 1)                                                          \n",
      "                                                                                                  \n",
      " conv_blocks_context/0/blocks/0  (None, None, None,   896        ['zero_padding3d[0][0]']         \n",
      " /conv (Conv3D)                 None, 32)                                                         \n",
      "                                                                                                  \n",
      " conv_blocks_context/0/blocks/0  (None, None, None,   64         ['conv_blocks_context/0/blocks/0/\n",
      " /norm (InstanceNormalization)  None, 32)                        conv[0][0]']                     \n",
      "                                                                                                  \n",
      " leaky_re_lu (LeakyReLU)        (None, None, None,   0           ['conv_blocks_context/0/blocks/0/\n",
      "                                None, 32)                        norm[0][0]']                     \n",
      "                                                                                                  \n",
      " dropout (Dropout)              (None, None, None,   0           ['leaky_re_lu[0][0]']            \n",
      "                                None, 32)                                                         \n",
      "                                                                                                  \n",
      " zero_padding3d_1 (ZeroPadding3  (None, None, None,   0          ['dropout[0][0]']                \n",
      " D)                             None, 32)                                                         \n",
      "                                                                                                  \n",
      " conv_blocks_context/0/blocks/1  (None, None, None,   27680      ['zero_padding3d_1[0][0]']       \n",
      " /conv (Conv3D)                 None, 32)                                                         \n",
      "                                                                                                  \n",
      " conv_blocks_context/0/blocks/1  (None, None, None,   64         ['conv_blocks_context/0/blocks/1/\n",
      " /norm (InstanceNormalization)  None, 32)                        conv[0][0]']                     \n",
      "                                                                                                  \n",
      " leaky_re_lu_1 (LeakyReLU)      (None, None, None,   0           ['conv_blocks_context/0/blocks/1/\n",
      "                                None, 32)                        norm[0][0]']                     \n",
      "                                                                                                  \n",
      " dropout_1 (Dropout)            (None, None, None,   0           ['leaky_re_lu_1[0][0]']          \n",
      "                                None, 32)                                                         \n",
      "                                                                                                  \n",
      " zero_padding3d_2 (ZeroPadding3  (None, None, None,   0          ['dropout_1[0][0]']              \n",
      " D)                             None, 32)                                                         \n",
      "                                                                                                  \n",
      " conv_blocks_context/1/blocks/0  (None, None, None,   55360      ['zero_padding3d_2[0][0]']       \n",
      " /conv (Conv3D)                 None, 64)                                                         \n",
      "                                                                                                  \n",
      " conv_blocks_context/1/blocks/0  (None, None, None,   128        ['conv_blocks_context/1/blocks/0/\n",
      " /norm (InstanceNormalization)  None, 64)                        conv[0][0]']                     \n",
      "                                                                                                  \n",
      " leaky_re_lu_2 (LeakyReLU)      (None, None, None,   0           ['conv_blocks_context/1/blocks/0/\n",
      "                                None, 64)                        norm[0][0]']                     \n",
      "                                                                                                  \n",
      " dropout_2 (Dropout)            (None, None, None,   0           ['leaky_re_lu_2[0][0]']          \n",
      "                                None, 64)                                                         \n",
      "                                                                                                  \n",
      " zero_padding3d_3 (ZeroPadding3  (None, None, None,   0          ['dropout_2[0][0]']              \n",
      " D)                             None, 64)                                                         \n",
      "                                                                                                  \n",
      " conv_blocks_context/1/blocks/1  (None, None, None,   110656     ['zero_padding3d_3[0][0]']       \n",
      " /conv (Conv3D)                 None, 64)                                                         \n",
      "                                                                                                  \n",
      " conv_blocks_context/1/blocks/1  (None, None, None,   128        ['conv_blocks_context/1/blocks/1/\n",
      " /norm (InstanceNormalization)  None, 64)                        conv[0][0]']                     \n",
      "                                                                                                  \n",
      " leaky_re_lu_3 (LeakyReLU)      (None, None, None,   0           ['conv_blocks_context/1/blocks/1/\n",
      "                                None, 64)                        norm[0][0]']                     \n",
      "                                                                                                  \n",
      " dropout_3 (Dropout)            (None, None, None,   0           ['leaky_re_lu_3[0][0]']          \n",
      "                                None, 64)                                                         \n",
      "                                                                                                  \n",
      " zero_padding3d_4 (ZeroPadding3  (None, None, None,   0          ['dropout_3[0][0]']              \n",
      " D)                             None, 64)                                                         \n",
      "                                                                                                  \n",
      " conv_blocks_context/2/blocks/0  (None, None, None,   221312     ['zero_padding3d_4[0][0]']       \n",
      " /conv (Conv3D)                 None, 128)                                                        \n",
      "                                                                                                  \n",
      " conv_blocks_context/2/blocks/0  (None, None, None,   256        ['conv_blocks_context/2/blocks/0/\n",
      " /norm (InstanceNormalization)  None, 128)                       conv[0][0]']                     \n",
      "                                                                                                  \n",
      " leaky_re_lu_4 (LeakyReLU)      (None, None, None,   0           ['conv_blocks_context/2/blocks/0/\n",
      "                                None, 128)                       norm[0][0]']                     \n",
      "                                                                                                  \n",
      " dropout_4 (Dropout)            (None, None, None,   0           ['leaky_re_lu_4[0][0]']          \n",
      "                                None, 128)                                                        \n",
      "                                                                                                  \n",
      " zero_padding3d_5 (ZeroPadding3  (None, None, None,   0          ['dropout_4[0][0]']              \n",
      " D)                             None, 128)                                                        \n",
      "                                                                                                  \n",
      " conv_blocks_context/2/blocks/1  (None, None, None,   442496     ['zero_padding3d_5[0][0]']       \n",
      " /conv (Conv3D)                 None, 128)                                                        \n",
      "                                                                                                  \n",
      " conv_blocks_context/2/blocks/1  (None, None, None,   256        ['conv_blocks_context/2/blocks/1/\n",
      " /norm (InstanceNormalization)  None, 128)                       conv[0][0]']                     \n",
      "                                                                                                  \n",
      " leaky_re_lu_5 (LeakyReLU)      (None, None, None,   0           ['conv_blocks_context/2/blocks/1/\n",
      "                                None, 128)                       norm[0][0]']                     \n",
      "                                                                                                  \n",
      " dropout_5 (Dropout)            (None, None, None,   0           ['leaky_re_lu_5[0][0]']          \n",
      "                                None, 128)                                                        \n",
      "                                                                                                  \n",
      " zero_padding3d_6 (ZeroPadding3  (None, None, None,   0          ['dropout_5[0][0]']              \n",
      " D)                             None, 128)                                                        \n",
      "                                                                                                  \n",
      " conv_blocks_context/3/blocks/0  (None, None, None,   884992     ['zero_padding3d_6[0][0]']       \n",
      " /conv (Conv3D)                 None, 256)                                                        \n",
      "                                                                                                  \n",
      " conv_blocks_context/3/blocks/0  (None, None, None,   512        ['conv_blocks_context/3/blocks/0/\n",
      " /norm (InstanceNormalization)  None, 256)                       conv[0][0]']                     \n",
      "                                                                                                  \n",
      " leaky_re_lu_6 (LeakyReLU)      (None, None, None,   0           ['conv_blocks_context/3/blocks/0/\n",
      "                                None, 256)                       norm[0][0]']                     \n",
      "                                                                                                  \n",
      " dropout_6 (Dropout)            (None, None, None,   0           ['leaky_re_lu_6[0][0]']          \n",
      "                                None, 256)                                                        \n",
      "                                                                                                  \n",
      " zero_padding3d_7 (ZeroPadding3  (None, None, None,   0          ['dropout_6[0][0]']              \n",
      " D)                             None, 256)                                                        \n",
      "                                                                                                  \n",
      " conv_blocks_context/3/blocks/1  (None, None, None,   1769728    ['zero_padding3d_7[0][0]']       \n",
      " /conv (Conv3D)                 None, 256)                                                        \n",
      "                                                                                                  \n",
      " conv_blocks_context/3/blocks/1  (None, None, None,   512        ['conv_blocks_context/3/blocks/1/\n",
      " /norm (InstanceNormalization)  None, 256)                       conv[0][0]']                     \n",
      "                                                                                                  \n",
      " leaky_re_lu_7 (LeakyReLU)      (None, None, None,   0           ['conv_blocks_context/3/blocks/1/\n",
      "                                None, 256)                       norm[0][0]']                     \n",
      "                                                                                                  \n",
      " dropout_7 (Dropout)            (None, None, None,   0           ['leaky_re_lu_7[0][0]']          \n",
      "                                None, 256)                                                        \n",
      "                                                                                                  \n",
      " zero_padding3d_8 (ZeroPadding3  (None, None, None,   0          ['dropout_7[0][0]']              \n",
      " D)                             None, 256)                                                        \n",
      "                                                                                                  \n",
      " conv_blocks_context/4/blocks/0  (None, None, None,   2212160    ['zero_padding3d_8[0][0]']       \n",
      " /conv (Conv3D)                 None, 320)                                                        \n",
      "                                                                                                  \n",
      " conv_blocks_context/4/blocks/0  (None, None, None,   640        ['conv_blocks_context/4/blocks/0/\n",
      " /norm (InstanceNormalization)  None, 320)                       conv[0][0]']                     \n",
      "                                                                                                  \n",
      " leaky_re_lu_8 (LeakyReLU)      (None, None, None,   0           ['conv_blocks_context/4/blocks/0/\n",
      "                                None, 320)                       norm[0][0]']                     \n",
      "                                                                                                  \n",
      " dropout_8 (Dropout)            (None, None, None,   0           ['leaky_re_lu_8[0][0]']          \n",
      "                                None, 320)                                                        \n",
      "                                                                                                  \n",
      " zero_padding3d_9 (ZeroPadding3  (None, None, None,   0          ['dropout_8[0][0]']              \n",
      " D)                             None, 320)                                                        \n",
      "                                                                                                  \n",
      " conv_blocks_context/4/blocks/1  (None, None, None,   2765120    ['zero_padding3d_9[0][0]']       \n",
      " /conv (Conv3D)                 None, 320)                                                        \n",
      "                                                                                                  \n",
      " conv_blocks_context/4/blocks/1  (None, None, None,   640        ['conv_blocks_context/4/blocks/1/\n",
      " /norm (InstanceNormalization)  None, 320)                       conv[0][0]']                     \n",
      "                                                                                                  \n",
      " leaky_re_lu_9 (LeakyReLU)      (None, None, None,   0           ['conv_blocks_context/4/blocks/1/\n",
      "                                None, 320)                       norm[0][0]']                     \n",
      "                                                                                                  \n",
      " dropout_9 (Dropout)            (None, None, None,   0           ['leaky_re_lu_9[0][0]']          \n",
      "                                None, 320)                                                        \n",
      "                                                                                                  \n",
      " zero_padding3d_10 (ZeroPadding  (None, None, None,   0          ['dropout_9[0][0]']              \n",
      " 3D)                            None, 320)                                                        \n",
      "                                                                                                  \n",
      " conv_blocks_context/5/0/blocks  (None, None, None,   2765120    ['zero_padding3d_10[0][0]']      \n",
      " /0/conv (Conv3D)               None, 320)                                                        \n",
      "                                                                                                  \n",
      " conv_blocks_context/5/0/blocks  (None, None, None,   640        ['conv_blocks_context/5/0/blocks/\n",
      " /0/norm (InstanceNormalization  None, 320)                      0/conv[0][0]']                   \n",
      " )                                                                                                \n",
      "                                                                                                  \n",
      " leaky_re_lu_10 (LeakyReLU)     (None, None, None,   0           ['conv_blocks_context/5/0/blocks/\n",
      "                                None, 320)                       0/norm[0][0]']                   \n",
      "                                                                                                  \n",
      " dropout_10 (Dropout)           (None, None, None,   0           ['leaky_re_lu_10[0][0]']         \n",
      "                                None, 320)                                                        \n",
      "                                                                                                  \n",
      " zero_padding3d_11 (ZeroPadding  (None, None, None,   0          ['dropout_10[0][0]']             \n",
      " 3D)                            None, 320)                                                        \n",
      "                                                                                                  \n",
      " conv_blocks_context/5/1/blocks  (None, None, None,   2765120    ['zero_padding3d_11[0][0]']      \n",
      " /1/conv (Conv3D)               None, 320)                                                        \n",
      "                                                                                                  \n",
      " conv_blocks_context/5/1/blocks  (None, None, None,   640        ['conv_blocks_context/5/1/blocks/\n",
      " /1/norm (InstanceNormalization  None, 320)                      1/conv[0][0]']                   \n",
      " )                                                                                                \n",
      "                                                                                                  \n",
      " leaky_re_lu_11 (LeakyReLU)     (None, None, None,   0           ['conv_blocks_context/5/1/blocks/\n",
      "                                None, 320)                       1/norm[0][0]']                   \n",
      "                                                                                                  \n",
      " dropout_11 (Dropout)           (None, None, None,   0           ['leaky_re_lu_11[0][0]']         \n",
      "                                None, 320)                                                        \n",
      "                                                                                                  \n",
      " tu/0 (Conv3DTranspose)         (None, None, None,   102400      ['dropout_11[0][0]']             \n",
      "                                None, 320)                                                        \n",
      "                                                                                                  \n",
      " concatenate (Concatenate)      (None, None, None,   0           ['tu/0[0][0]',                   \n",
      "                                None, 640)                        'dropout_9[0][0]']              \n",
      "                                                                                                  \n",
      " zero_padding3d_12 (ZeroPadding  (None, None, None,   0          ['concatenate[0][0]']            \n",
      " 3D)                            None, 640)                                                        \n",
      "                                                                                                  \n",
      " conv_blocks_localization/0/0/b  (None, None, None,   5529920    ['zero_padding3d_12[0][0]']      \n",
      " locks/0/conv (Conv3D)          None, 320)                                                        \n",
      "                                                                                                  \n",
      " conv_blocks_localization/0/0/b  (None, None, None,   640        ['conv_blocks_localization/0/0/bl\n",
      " locks/0/norm (InstanceNormaliz  None, 320)                      ocks/0/conv[0][0]']              \n",
      " ation)                                                                                           \n",
      "                                                                                                  \n",
      " leaky_re_lu_12 (LeakyReLU)     (None, None, None,   0           ['conv_blocks_localization/0/0/bl\n",
      "                                None, 320)                       ocks/0/norm[0][0]']              \n",
      "                                                                                                  \n",
      " dropout_12 (Dropout)           (None, None, None,   0           ['leaky_re_lu_12[0][0]']         \n",
      "                                None, 320)                                                        \n",
      "                                                                                                  \n",
      " zero_padding3d_13 (ZeroPadding  (None, None, None,   0          ['dropout_12[0][0]']             \n",
      " 3D)                            None, 320)                                                        \n",
      "                                                                                                  \n",
      " conv_blocks_localization/0/1/b  (None, None, None,   2765120    ['zero_padding3d_13[0][0]']      \n",
      " locks/1/conv (Conv3D)          None, 320)                                                        \n",
      "                                                                                                  \n",
      " conv_blocks_localization/0/1/b  (None, None, None,   640        ['conv_blocks_localization/0/1/bl\n",
      " locks/1/norm (InstanceNormaliz  None, 320)                      ocks/1/conv[0][0]']              \n",
      " ation)                                                                                           \n",
      "                                                                                                  \n",
      " leaky_re_lu_13 (LeakyReLU)     (None, None, None,   0           ['conv_blocks_localization/0/1/bl\n",
      "                                None, 320)                       ocks/1/norm[0][0]']              \n",
      "                                                                                                  \n",
      " dropout_13 (Dropout)           (None, None, None,   0           ['leaky_re_lu_13[0][0]']         \n",
      "                                None, 320)                                                        \n",
      "                                                                                                  \n",
      " tu/1 (Conv3DTranspose)         (None, None, None,   655360      ['dropout_13[0][0]']             \n",
      "                                None, 256)                                                        \n",
      "                                                                                                  \n",
      " concatenate_1 (Concatenate)    (None, None, None,   0           ['tu/1[0][0]',                   \n",
      "                                None, 512)                        'dropout_7[0][0]']              \n",
      "                                                                                                  \n",
      " zero_padding3d_14 (ZeroPadding  (None, None, None,   0          ['concatenate_1[0][0]']          \n",
      " 3D)                            None, 512)                                                        \n",
      "                                                                                                  \n",
      " conv_blocks_localization/1/0/b  (None, None, None,   3539200    ['zero_padding3d_14[0][0]']      \n",
      " locks/0/conv (Conv3D)          None, 256)                                                        \n",
      "                                                                                                  \n",
      " conv_blocks_localization/1/0/b  (None, None, None,   512        ['conv_blocks_localization/1/0/bl\n",
      " locks/0/norm (InstanceNormaliz  None, 256)                      ocks/0/conv[0][0]']              \n",
      " ation)                                                                                           \n",
      "                                                                                                  \n",
      " leaky_re_lu_14 (LeakyReLU)     (None, None, None,   0           ['conv_blocks_localization/1/0/bl\n",
      "                                None, 256)                       ocks/0/norm[0][0]']              \n",
      "                                                                                                  \n",
      " dropout_14 (Dropout)           (None, None, None,   0           ['leaky_re_lu_14[0][0]']         \n",
      "                                None, 256)                                                        \n",
      "                                                                                                  \n",
      " zero_padding3d_15 (ZeroPadding  (None, None, None,   0          ['dropout_14[0][0]']             \n",
      " 3D)                            None, 256)                                                        \n",
      "                                                                                                  \n",
      " conv_blocks_localization/1/1/b  (None, None, None,   1769728    ['zero_padding3d_15[0][0]']      \n",
      " locks/1/conv (Conv3D)          None, 256)                                                        \n",
      "                                                                                                  \n",
      " conv_blocks_localization/1/1/b  (None, None, None,   512        ['conv_blocks_localization/1/1/bl\n",
      " locks/1/norm (InstanceNormaliz  None, 256)                      ocks/1/conv[0][0]']              \n",
      " ation)                                                                                           \n",
      "                                                                                                  \n",
      " leaky_re_lu_15 (LeakyReLU)     (None, None, None,   0           ['conv_blocks_localization/1/1/bl\n",
      "                                None, 256)                       ocks/1/norm[0][0]']              \n",
      "                                                                                                  \n",
      " dropout_15 (Dropout)           (None, None, None,   0           ['leaky_re_lu_15[0][0]']         \n",
      "                                None, 256)                                                        \n",
      "                                                                                                  \n",
      " tu/2 (Conv3DTranspose)         (None, None, None,   262144      ['dropout_15[0][0]']             \n",
      "                                None, 128)                                                        \n",
      "                                                                                                  \n",
      " concatenate_2 (Concatenate)    (None, None, None,   0           ['tu/2[0][0]',                   \n",
      "                                None, 256)                        'dropout_5[0][0]']              \n",
      "                                                                                                  \n",
      " zero_padding3d_16 (ZeroPadding  (None, None, None,   0          ['concatenate_2[0][0]']          \n",
      " 3D)                            None, 256)                                                        \n",
      "                                                                                                  \n",
      " conv_blocks_localization/2/0/b  (None, None, None,   884864     ['zero_padding3d_16[0][0]']      \n",
      " locks/0/conv (Conv3D)          None, 128)                                                        \n",
      "                                                                                                  \n",
      " conv_blocks_localization/2/0/b  (None, None, None,   256        ['conv_blocks_localization/2/0/bl\n",
      " locks/0/norm (InstanceNormaliz  None, 128)                      ocks/0/conv[0][0]']              \n",
      " ation)                                                                                           \n",
      "                                                                                                  \n",
      " leaky_re_lu_16 (LeakyReLU)     (None, None, None,   0           ['conv_blocks_localization/2/0/bl\n",
      "                                None, 128)                       ocks/0/norm[0][0]']              \n",
      "                                                                                                  \n",
      " dropout_16 (Dropout)           (None, None, None,   0           ['leaky_re_lu_16[0][0]']         \n",
      "                                None, 128)                                                        \n",
      "                                                                                                  \n",
      " zero_padding3d_17 (ZeroPadding  (None, None, None,   0          ['dropout_16[0][0]']             \n",
      " 3D)                            None, 128)                                                        \n",
      "                                                                                                  \n",
      " conv_blocks_localization/2/1/b  (None, None, None,   442496     ['zero_padding3d_17[0][0]']      \n",
      " locks/1/conv (Conv3D)          None, 128)                                                        \n",
      "                                                                                                  \n",
      " conv_blocks_localization/2/1/b  (None, None, None,   256        ['conv_blocks_localization/2/1/bl\n",
      " locks/1/norm (InstanceNormaliz  None, 128)                      ocks/1/conv[0][0]']              \n",
      " ation)                                                                                           \n",
      "                                                                                                  \n",
      " leaky_re_lu_17 (LeakyReLU)     (None, None, None,   0           ['conv_blocks_localization/2/1/bl\n",
      "                                None, 128)                       ocks/1/norm[0][0]']              \n",
      "                                                                                                  \n",
      " dropout_17 (Dropout)           (None, None, None,   0           ['leaky_re_lu_17[0][0]']         \n",
      "                                None, 128)                                                        \n",
      "                                                                                                  \n",
      " tu/3 (Conv3DTranspose)         (None, None, None,   65536       ['dropout_17[0][0]']             \n",
      "                                None, 64)                                                         \n",
      "                                                                                                  \n",
      " concatenate_3 (Concatenate)    (None, None, None,   0           ['tu/3[0][0]',                   \n",
      "                                None, 128)                        'dropout_3[0][0]']              \n",
      "                                                                                                  \n",
      " zero_padding3d_18 (ZeroPadding  (None, None, None,   0          ['concatenate_3[0][0]']          \n",
      " 3D)                            None, 128)                                                        \n",
      "                                                                                                  \n",
      " conv_blocks_localization/3/0/b  (None, None, None,   221248     ['zero_padding3d_18[0][0]']      \n",
      " locks/0/conv (Conv3D)          None, 64)                                                         \n",
      "                                                                                                  \n",
      " conv_blocks_localization/3/0/b  (None, None, None,   128        ['conv_blocks_localization/3/0/bl\n",
      " locks/0/norm (InstanceNormaliz  None, 64)                       ocks/0/conv[0][0]']              \n",
      " ation)                                                                                           \n",
      "                                                                                                  \n",
      " leaky_re_lu_18 (LeakyReLU)     (None, None, None,   0           ['conv_blocks_localization/3/0/bl\n",
      "                                None, 64)                        ocks/0/norm[0][0]']              \n",
      "                                                                                                  \n",
      " dropout_18 (Dropout)           (None, None, None,   0           ['leaky_re_lu_18[0][0]']         \n",
      "                                None, 64)                                                         \n",
      "                                                                                                  \n",
      " zero_padding3d_19 (ZeroPadding  (None, None, None,   0          ['dropout_18[0][0]']             \n",
      " 3D)                            None, 64)                                                         \n",
      "                                                                                                  \n",
      " conv_blocks_localization/3/1/b  (None, None, None,   110656     ['zero_padding3d_19[0][0]']      \n",
      " locks/1/conv (Conv3D)          None, 64)                                                         \n",
      "                                                                                                  \n",
      " conv_blocks_localization/3/1/b  (None, None, None,   128        ['conv_blocks_localization/3/1/bl\n",
      " locks/1/norm (InstanceNormaliz  None, 64)                       ocks/1/conv[0][0]']              \n",
      " ation)                                                                                           \n",
      "                                                                                                  \n",
      " leaky_re_lu_19 (LeakyReLU)     (None, None, None,   0           ['conv_blocks_localization/3/1/bl\n",
      "                                None, 64)                        ocks/1/norm[0][0]']              \n",
      "                                                                                                  \n",
      " dropout_19 (Dropout)           (None, None, None,   0           ['leaky_re_lu_19[0][0]']         \n",
      "                                None, 64)                                                         \n",
      "                                                                                                  \n",
      " tu/4 (Conv3DTranspose)         (None, None, None,   16384       ['dropout_19[0][0]']             \n",
      "                                None, 32)                                                         \n",
      "                                                                                                  \n",
      " concatenate_4 (Concatenate)    (None, None, None,   0           ['tu/4[0][0]',                   \n",
      "                                None, 64)                         'dropout_1[0][0]']              \n",
      "                                                                                                  \n",
      " zero_padding3d_20 (ZeroPadding  (None, None, None,   0          ['concatenate_4[0][0]']          \n",
      " 3D)                            None, 64)                                                         \n",
      "                                                                                                  \n",
      " conv_blocks_localization/4/0/b  (None, None, None,   55328      ['zero_padding3d_20[0][0]']      \n",
      " locks/0/conv (Conv3D)          None, 32)                                                         \n",
      "                                                                                                  \n",
      " conv_blocks_localization/4/0/b  (None, None, None,   64         ['conv_blocks_localization/4/0/bl\n",
      " locks/0/norm (InstanceNormaliz  None, 32)                       ocks/0/conv[0][0]']              \n",
      " ation)                                                                                           \n",
      "                                                                                                  \n",
      " leaky_re_lu_20 (LeakyReLU)     (None, None, None,   0           ['conv_blocks_localization/4/0/bl\n",
      "                                None, 32)                        ocks/0/norm[0][0]']              \n",
      "                                                                                                  \n",
      " dropout_20 (Dropout)           (None, None, None,   0           ['leaky_re_lu_20[0][0]']         \n",
      "                                None, 32)                                                         \n",
      "                                                                                                  \n",
      " zero_padding3d_21 (ZeroPadding  (None, None, None,   0          ['dropout_20[0][0]']             \n",
      " 3D)                            None, 32)                                                         \n",
      "                                                                                                  \n",
      " conv_blocks_localization/4/1/b  (None, None, None,   27680      ['zero_padding3d_21[0][0]']      \n",
      " locks/1/conv (Conv3D)          None, 32)                                                         \n",
      "                                                                                                  \n",
      " conv_blocks_localization/4/1/b  (None, None, None,   64         ['conv_blocks_localization/4/1/bl\n",
      " locks/1/norm (InstanceNormaliz  None, 32)                       ocks/1/conv[0][0]']              \n",
      " ation)                                                                                           \n",
      "                                                                                                  \n",
      " leaky_re_lu_21 (LeakyReLU)     (None, None, None,   0           ['conv_blocks_localization/4/1/bl\n",
      "                                None, 32)                        ocks/1/norm[0][0]']              \n",
      "                                                                                                  \n",
      " dropout_21 (Dropout)           (None, None, None,   0           ['leaky_re_lu_21[0][0]']         \n",
      "                                None, 32)                                                         \n",
      "                                                                                                  \n",
      " seg_outputs/4 (Conv3D)         (None, None, None,   3360        ['dropout_21[0][0]']             \n",
      "                                None, 105)                                                        \n",
      "                                                                                                  \n",
      " activation (Activation)        (None, None, None,   0           ['seg_outputs/4[0][0]']          \n",
      "                                None, 105)                                                        \n",
      "                                                                                                  \n",
      "==================================================================================================\n",
      "Total params: 30,479,744\n",
      "Trainable params: 30,479,744\n",
      "Non-trainable params: 0\n",
      "__________________________________________________________________________________________________\n"
     ]
    }
   ],
   "source": [
    "model.summary()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e136a73e-e0b2-4d22-955b-9fa3cecd0af8",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "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.13"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}