{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "d422208b",
   "metadata": {},
   "outputs": [],
   "source": [
    "import autoreload\n",
    "%load_ext autoreload\n",
    "%autoreload 2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "2785de6a",
   "metadata": {},
   "outputs": [],
   "source": [
    "%reload_ext autoreload"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "934aac57",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "os.environ[\"CUDA_DEVICE_ORDER\"] = \"PCI_BUS_ID\"   # see issue #152\n",
    "os.environ[\"CUDA_VISIBLE_DEVICES\"] = \"\"\n",
    "os.environ[\"TF_KERAS\"] = \"1\"\n",
    "# os.environ['TF_CPP_MIN_LOG_LEVEL'] = '2'  "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "c677106c",
   "metadata": {},
   "outputs": [],
   "source": [
    "import tensorflow as tf\n",
    "import keras\n",
    "from tensorflow.keras.models import Sequential, Model\n",
    "from tensorflow.keras.layers import Input , Dense, Lambda, Normalization, Concatenate\n",
    "from tensorflow.keras.activations import relu, sigmoid, tanh\n",
    "from tensorflow.keras import backend as K\n",
    "from tensorflow.keras.optimizers import Adam, SGD, schedules\n",
    "from tensorflow_addons.optimizers import AdamW, extend_with_decoupled_weight_decay\n",
    "\n",
    "from tensorflow.keras.callbacks import ModelCheckpoint, TerminateOnNaN, TensorBoard, \\\n",
    "                                        EarlyStopping, ReduceLROnPlateau, CSVLogger, Callback\n",
    "\n",
    "import numpy as np\n",
    "import random\n",
    "import os\n",
    "\n",
    "import gc\n",
    "\n",
    "from tqdm import tqdm "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "a3141967",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "importing Jupyter notebook from modules.ipynb\n",
      "importing Jupyter notebook from loss.ipynb\n",
      "importing Jupyter notebook from bandERB.ipynb\n",
      "importing Jupyter notebook from params.ipynb\n",
      "importing Jupyter notebook from utils.ipynb\n",
      "importing Jupyter notebook from dataloader.ipynb\n"
     ]
    }
   ],
   "source": [
    "import import_ipynb\n",
    "\n",
    "from modules import convkxf, GroupFC#, GroupGRULayer\n",
    "\n",
    "from loss import MaskLoss, LocalSnrTarget, DfAlphaLoss, SpectralLoss\n",
    "from utils import mask_operations, df_operations, synthesis_frame, df_operations_fast\n",
    "from params import model_params\n",
    "from dataloader import read_tfrecod_data\n",
    "from bandERB import ERBBand, ERB_pro_matrix"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "31a22c7c",
   "metadata": {},
   "outputs": [],
   "source": [
    "stateful = False\n",
    "gru_norm, fc_norm = False, False\n",
    "if stateful: from modules import GroupGRULayer_stateful as GroupGRULayer\n",
    "else: from modules import GroupGRULayer_lite as GroupGRULayer"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "b38dff81",
   "metadata": {},
   "outputs": [],
   "source": [
    "# from tensorflow.compat.v1 import ConfigProto, InteractiveSession, enable_eager_execution\n",
    "\n",
    "# config = ConfigProto(allow_soft_placement=True)\n",
    "# config.gpu_options.per_process_gpu_memory_fraction = 0.6\n",
    "# config.gpu_options.allow_growth = True\n",
    "# InteractiveSession(config=config)\n",
    "# enable_eager_execution(config=config)\n",
    "\n",
    "# # tf.compat.v1.disable_eager_execution()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8f0e15be",
   "metadata": {},
   "source": [
    "# Read config"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "5bce3a7b",
   "metadata": {},
   "outputs": [],
   "source": [
    "p = model_params('config.ini')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "818436e6",
   "metadata": {},
   "outputs": [],
   "source": [
    "def as_complex(x):\n",
    "    if x.dtype == np.complex64 or x.dtype == np.complex128:\n",
    "        return x\n",
    "    else:\n",
    "        return np.complex(x[...,-2], x[..., -1])\n",
    "\n",
    "def as_real(x):\n",
    "    if x.dtype == np.complex64 or x.dtype == np.complex128:\n",
    "        return np.concatenate([np.expand_dims(np.real(x),axis=-1), \n",
    "                               np.expand_dims(np.imag(x),axis=-1)],axis=-1)\n",
    "    else: return x"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "7829fc8c",
   "metadata": {},
   "outputs": [],
   "source": [
    "gru_layer = 2\n",
    "p.fc_group = 8\n",
    "p.gru_group = 1"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "fd863e94",
   "metadata": {},
   "source": [
    "# Model construction"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "ed6f8283",
   "metadata": {},
   "outputs": [],
   "source": [
    "if stateful: time = None\n",
    "else: time = 3\n",
    "    \n",
    "erb_inputs = Input(shape=(time, p.nb_erb, 1), batch_size=1, name='ERB_input')\n",
    "spec_inputs = Input(shape=(time, p.nb_df, 2), batch_size=1, name='spec_input') # complex\n",
    "\n",
    "state = Input(shape=(256), batch_size=1, name='state0')\n",
    "\n",
    "state10 = Input(shape=(256), batch_size=1, name='state10')\n",
    "state11 = Input(shape=(256), batch_size=1, name='state11')\n",
    "state1 = [state10, state11]\n",
    "\n",
    "state20 = Input(shape=(256), batch_size=1, name='state20')\n",
    "state21 = Input(shape=(256), batch_size=1, name='state21')\n",
    "\n",
    "if gru_layer>1:\n",
    "    state1 = [state10, state11]\n",
    "    state2 = [state20, state21]\n",
    "else:\n",
    "    state1 = [state10]\n",
    "    state2 = [state20]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "61021b38",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'filters': 64, 'groups': 1, 'kernel_size': (3, 3), 'strides': (1, 1), 'use_bias': False, 'kernel_initializer': 'he_normal', 'padding': 'valid', 'name': 'conv0_encoder'}\n",
      "{'filters': 64, 'groups': 64, 'kernel_size': (1, 3), 'strides': (1, 2), 'use_bias': False, 'kernel_initializer': 'he_normal', 'padding': 'valid', 'name': 'conv1_encoder'}\n",
      "{'filters': 64, 'groups': 64, 'kernel_size': (1, 3), 'strides': (1, 2), 'use_bias': False, 'kernel_initializer': 'he_normal', 'padding': 'valid', 'name': 'conv2_encoder'}\n",
      "{'filters': 64, 'groups': 64, 'kernel_size': (1, 3), 'strides': (1, 1), 'use_bias': False, 'kernel_initializer': 'he_normal', 'padding': 'valid', 'name': 'conv3_encoder'}\n",
      "{'filters': 64, 'groups': 2, 'kernel_size': (3, 3), 'strides': (1, 1), 'use_bias': False, 'kernel_initializer': 'he_normal', 'padding': 'valid', 'name': 'df_conv0_encoder'}\n",
      "{'filters': 64, 'groups': 64, 'kernel_size': (1, 3), 'strides': (1, 2), 'use_bias': False, 'kernel_initializer': 'he_normal', 'padding': 'valid', 'name': 'df_conv1_encoder'}\n"
     ]
    }
   ],
   "source": [
    "erb_conv0 = convkxf(erb_inputs, p.conv_out_ch, k=3, f=3, fstride=1, bias=False, batch_norm=True, \n",
    "                    training=False, name='conv0_encoder')\n",
    "erb_conv1 = convkxf(erb_conv0, p.conv_out_ch, k=1, f=3, fstride=2, bias=False, batch_norm=True, \n",
    "                     name='conv1_encoder')\n",
    "erb_conv2 = convkxf(erb_conv1, p.conv_out_ch, k=1, f=3, fstride=2, bias=False, batch_norm=True, \n",
    "                     name='conv2_encoder')\n",
    "erb_conv3 = convkxf(erb_conv2, p.conv_out_ch, k=1, f=3, fstride=1, bias=False, batch_norm=True, \n",
    "                     name='conv3_encoder')\n",
    "\n",
    "df_conv0 = convkxf(spec_inputs, p.conv_out_ch, k=3, f=3, fstride=1, batch_norm=True, \n",
    "                   training=False,  name='df_conv0_encoder')\n",
    "df_conv1 = convkxf(df_conv0, p.conv_out_ch, k=1, f=3, fstride=2, batch_norm=True, \n",
    "                    name='df_conv1_encoder')\n",
    "\n",
    "shape = [tf.shape(df_conv1)[l] for l in range(4)]\n",
    "df_conv1 = tf.transpose(df_conv1,(0,1,3,2))\n",
    "df_conv1 = tf.reshape(df_conv1, [shape[0], shape[1], shape[2]*shape[3]])\n",
    "\n",
    "cemb = GroupFC(df_conv1, p.fc_hidden, p.fc_group, name='GFC_encoder')\n",
    "    \n",
    "shape = [tf.shape(erb_conv3)[l] for l in range(4)]  \n",
    "emb = Concatenate()([tf.reshape(tf.transpose(erb_conv3,(0,1,3,2)), [shape[0], shape[1], shape[2]*shape[3]]), cemb])\n",
    "\n",
    "if p.model_ver == \"3.0\":\n",
    "    emb = GroupFC(emb, p.gru_hidden, p.fc_group, name='GFC_emb')\n",
    "    \n",
    "if stateful:\n",
    "    GRU_emb = GroupGRULayer(emb, p.gru_hidden, p.gru_group, num_layer=1, name='GGRU0', \n",
    "                               add_output=True, training=False, norm=gru_norm)\n",
    "else:\n",
    "\n",
    "    GRU_emb, lstate_o = GroupGRULayer(emb, p.gru_hidden, p.gru_group, num_layer=1, state=[state], name='GGRU0', \n",
    "                               add_output=True, training=False, norm=gru_norm)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "6572e9ae",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'filters': 64, 'groups': 64, 'kernel_size': (1, 1), 'strides': 1, 'use_bias': False, 'kernel_initializer': 'he_normal', 'padding': 'valid', 'name': 'convp3'}\n",
      "{'filters': 64, 'groups': 64, 'kernel_size': (1, 3), 'strides': (1, 1), 'use_bias': False, 'kernel_initializer': 'he_normal', 'padding': 'valid', 'name': 'convt3'}\n",
      "{'filters': 64, 'groups': 64, 'kernel_size': (1, 1), 'strides': 1, 'use_bias': False, 'kernel_initializer': 'he_normal', 'padding': 'valid', 'name': 'convp2'}\n",
      "{'filters': 64, 'groups': 64, 'kernel_size': (1, 1), 'strides': 1, 'use_bias': False, 'kernel_initializer': 'he_normal', 'padding': 'valid', 'name': 'convp1'}\n",
      "{'filters': 64, 'groups': 64, 'kernel_size': (1, 1), 'strides': 1, 'use_bias': False, 'kernel_initializer': 'he_normal', 'padding': 'valid', 'name': 'convp0'}\n",
      "{'filters': 1, 'groups': 1, 'kernel_size': (1, 3), 'strides': (1, 1), 'use_bias': False, 'kernel_initializer': 'he_normal', 'padding': 'valid', 'name': 'mask_out'}\n"
     ]
    }
   ],
   "source": [
    "if stateful:\n",
    "    GRU_emb1 = GroupGRULayer(GRU_emb, p.gru_hidden, p.gru_group, num_layer=gru_layer, name='GGRU01', \n",
    "                               add_output=True, training=False, norm=gru_norm)\n",
    "else:\n",
    "    GRU_emb1, lstate1_o = GroupGRULayer(GRU_emb, p.gru_hidden, p.gru_group, num_layer=gru_layer, state=state1, name='GGRU01', \n",
    "                                 add_output=True, training=False, norm=gru_norm)\n",
    "de_emb = GroupFC(GRU_emb1, p.fc_hidden, p.fc_group, activation='relu', \n",
    "                 name='GFC_decoder', norm=fc_norm)\n",
    "\n",
    "shape = [tf.shape(erb_conv3)[l] for l in range(4)]\n",
    "emb_decoder = tf.reshape(de_emb, shape = [shape[0], shape[1], shape[2], -1], name='decoder_reshape')\n",
    "\n",
    "kwargs = {\n",
    "    \"k\": 1,\n",
    "    \"batch_norm\": True,\n",
    "}\n",
    "tkwargs = {\n",
    "    \"k\": 1,\n",
    "    \"batch_norm\": True,\n",
    "    \"mode\": \"transposed\"\n",
    "}\n",
    "pkwargs = {\n",
    "    \"k\": 1,\n",
    "    \"f\": 1,\n",
    "    \"batch_norm\": True\n",
    "}\n",
    "\n",
    "convp3 = convkxf(erb_conv3, out_ch=p.conv_out_ch, name='convp3', **pkwargs)  # Conv\n",
    "vt3_in = convp3 + emb_decoder\n",
    "convt3 = convkxf(vt3_in, out_ch=p.conv_out_ch, fstride=1, name='convt3', **kwargs) #ConvT\n",
    "\n",
    "convp2 = convkxf(erb_conv2, out_ch=p.conv_out_ch, name='convp2', **pkwargs) # Conv\n",
    "vt2_in = convp2 + convt3\n",
    "convt2 = convkxf(vt2_in, out_ch=p.conv_out_ch, name='convt2', **tkwargs)\n",
    "\n",
    "convp1 = convkxf(erb_conv1, out_ch=p.conv_out_ch, name='convp1', **pkwargs) # Conv\n",
    "vt1_in = convp1 + convt2\n",
    "convt1 = convkxf(vt1_in, out_ch=p.conv_out_ch, name='convt1', **tkwargs)\n",
    "\n",
    "convp0 = convkxf(erb_conv0, out_ch=p.conv_out_ch, name='convp0', **pkwargs) # Conv\n",
    "vt0_in = convp0 + convt1\n",
    "mask_out = convkxf(vt0_in, out_ch=1, k=1, fstride=1, act='sigmoid', batch_norm=False, \n",
    "                   reshape=True, name='mask_out')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "id": "18193a0f",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'filters': 10, 'groups': 2, 'kernel_size': (1, 1), 'strides': 1, 'use_bias': False, 'kernel_initializer': 'he_normal', 'padding': 'valid', 'name': 'convp_DfDecoder'}\n"
     ]
    }
   ],
   "source": [
    "if stateful:\n",
    "    GRU_emb2 = GroupGRULayer(GRU_emb, p.gru_hidden, p.gru_group, num_layer=gru_layer, name='GGRU1', \n",
    "                               add_output=True, training=False, norm=gru_norm)\n",
    "else:\n",
    "    GRU_emb2, lstate2_o = GroupGRULayer(GRU_emb, p.gru_hidden, p.gru_group, num_layer=gru_layer, state=state2, name='GGRU1', \n",
    "                                 add_output=True, training=False, norm=gru_norm)\n",
    "convp = convkxf(df_conv0, 2*p.df_order, k=1, f=1, complex_in=True, batch_norm=True, name='convp_DfDecoder') \n",
    "convp = tf.transpose(convp, (0,1,3,2))\n",
    "\n",
    "df_alpha = Dense(1, name='convp_alpha', activation='sigmoid')(GRU_emb2)\n",
    "\n",
    "c = GroupFC(GRU_emb2, p.nb_df*p.df_order*2, p.fc_group, activation='tanh', name='GFC_c')\n",
    "\n",
    "shape = [tf.shape(c)[k] for k in range(2)]\n",
    "c = tf.reshape(c,(shape[0], shape[1], p.df_order*2, p.nb_df))\n",
    "\n",
    "c = tf.reshape(c + convp,(shape[0], shape[1], p.df_order, 2, p.nb_df))\n",
    "\n",
    "df_coeff = tf.transpose(c, [0, 1, 2, 4, 3], name='output_c')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7b2d2b67",
   "metadata": {},
   "source": [
    "# Model compile"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "id": "635ef4ec",
   "metadata": {},
   "outputs": [],
   "source": [
    "if stateful:\n",
    "    inputs = [erb_inputs, spec_inputs]\n",
    "    outputs = [mask_out, df_alpha, df_coeff]\n",
    "else:\n",
    "    inputs = [erb_inputs, spec_inputs, state, state1, state2]\n",
    "    outputs = [mask_out, df_alpha, df_coeff, lstate_o, lstate1_o, lstate2_o]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "id": "a1fc6c10",
   "metadata": {},
   "outputs": [],
   "source": [
    "# inputs = [erb_inputs, spec_inputs, state, state1] #, state]\n",
    "# outputs = [mask_out, lstate_o, lstate1_o]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "id": "8f43f3ed",
   "metadata": {},
   "outputs": [],
   "source": [
    "model=Model(inputs=inputs, outputs=outputs, name='DfNet')\n",
    "model.trainable = False"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "id": "39435d21",
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Model: \"DfNet\"\n",
      "__________________________________________________________________________________________________\n",
      " Layer (type)                   Output Shape         Param #     Connected to                     \n",
      "==================================================================================================\n",
      " ERB_input (InputLayer)         [(1, 3, 32, 1)]      0           []                               \n",
      "                                                                                                  \n",
      " tf.compat.v1.pad (TFOpLambda)  (1, 3, 34, 1)        0           ['ERB_input[0][0]']              \n",
      "                                                                                                  \n",
      " spec_input (InputLayer)        [(1, 3, 96, 2)]      0           []                               \n",
      "                                                                                                  \n",
      " conv0_encoder (Conv2D)         (1, 1, 32, 64)       576         ['tf.compat.v1.pad[0][0]']       \n",
      "                                                                                                  \n",
      " tf.compat.v1.pad_4 (TFOpLambda  (1, 3, 98, 2)       0           ['spec_input[0][0]']             \n",
      " )                                                                                                \n",
      "                                                                                                  \n",
      " conv0_encoderBatchNorm (BatchN  (1, 1, 32, 64)      256         ['conv0_encoder[0][0]']          \n",
      " ormalization)                                                                                    \n",
      "                                                                                                  \n",
      " df_conv0_encoder (Conv2D)      (1, 1, 96, 64)       576         ['tf.compat.v1.pad_4[0][0]']     \n",
      "                                                                                                  \n",
      " tf.nn.relu (TFOpLambda)        (1, 1, 32, 64)       0           ['conv0_encoderBatchNorm[0][0]'] \n",
      "                                                                                                  \n",
      " df_conv0_encoder_1x1 (Conv2D)  (1, 1, 96, 64)       4096        ['df_conv0_encoder[0][0]']       \n",
      "                                                                                                  \n",
      " tf.compat.v1.pad_1 (TFOpLambda  (1, 1, 34, 64)      0           ['tf.nn.relu[0][0]']             \n",
      " )                                                                                                \n",
      "                                                                                                  \n",
      " df_conv0_encoderBatchNorm (Bat  (1, 1, 96, 64)      256         ['df_conv0_encoder_1x1[0][0]']   \n",
      " chNormalization)                                                                                 \n",
      "                                                                                                  \n",
      " conv1_encoder (Conv2D)         (1, 1, 16, 64)       192         ['tf.compat.v1.pad_1[0][0]']     \n",
      "                                                                                                  \n",
      " tf.nn.relu_4 (TFOpLambda)      (1, 1, 96, 64)       0           ['df_conv0_encoderBatchNorm[0][0]\n",
      "                                                                 ']                               \n",
      "                                                                                                  \n",
      " conv1_encoder_1x1 (Conv2D)     (1, 1, 16, 64)       4096        ['conv1_encoder[0][0]']          \n",
      "                                                                                                  \n",
      " tf.compat.v1.pad_5 (TFOpLambda  (1, 1, 98, 64)      0           ['tf.nn.relu_4[0][0]']           \n",
      " )                                                                                                \n",
      "                                                                                                  \n",
      " conv1_encoderBatchNorm (BatchN  (1, 1, 16, 64)      256         ['conv1_encoder_1x1[0][0]']      \n",
      " ormalization)                                                                                    \n",
      "                                                                                                  \n",
      " df_conv1_encoder (Conv2D)      (1, 1, 48, 64)       192         ['tf.compat.v1.pad_5[0][0]']     \n",
      "                                                                                                  \n",
      " tf.nn.relu_1 (TFOpLambda)      (1, 1, 16, 64)       0           ['conv1_encoderBatchNorm[0][0]'] \n",
      "                                                                                                  \n",
      " df_conv1_encoder_1x1 (Conv2D)  (1, 1, 48, 64)       4096        ['df_conv1_encoder[0][0]']       \n",
      "                                                                                                  \n",
      " tf.compat.v1.pad_2 (TFOpLambda  (1, 1, 18, 64)      0           ['tf.nn.relu_1[0][0]']           \n",
      " )                                                                                                \n",
      "                                                                                                  \n",
      " df_conv1_encoderBatchNorm (Bat  (1, 1, 48, 64)      256         ['df_conv1_encoder_1x1[0][0]']   \n",
      " chNormalization)                                                                                 \n",
      "                                                                                                  \n",
      " conv2_encoder (Conv2D)         (1, 1, 8, 64)        192         ['tf.compat.v1.pad_2[0][0]']     \n",
      "                                                                                                  \n",
      " tf.nn.relu_5 (TFOpLambda)      (1, 1, 48, 64)       0           ['df_conv1_encoderBatchNorm[0][0]\n",
      "                                                                 ']                               \n",
      "                                                                                                  \n",
      " conv2_encoder_1x1 (Conv2D)     (1, 1, 8, 64)        4096        ['conv2_encoder[0][0]']          \n",
      "                                                                                                  \n",
      " tf.compat.v1.shape_2 (TFOpLamb  (4,)                0           ['tf.nn.relu_5[0][0]']           \n",
      " da)                                                                                              \n",
      "                                                                                                  \n",
      " tf.compat.v1.shape_3 (TFOpLamb  (4,)                0           ['tf.nn.relu_5[0][0]']           \n",
      " da)                                                                                              \n",
      "                                                                                                  \n",
      " conv2_encoderBatchNorm (BatchN  (1, 1, 8, 64)       256         ['conv2_encoder_1x1[0][0]']      \n",
      " ormalization)                                                                                    \n",
      "                                                                                                  \n",
      " tf.compat.v1.shape (TFOpLambda  (4,)                0           ['tf.nn.relu_5[0][0]']           \n",
      " )                                                                                                \n",
      "                                                                                                  \n",
      " tf.compat.v1.shape_1 (TFOpLamb  (4,)                0           ['tf.nn.relu_5[0][0]']           \n",
      " da)                                                                                              \n",
      "                                                                                                  \n",
      " tf.__operators__.getitem_2 (Sl  ()                  0           ['tf.compat.v1.shape_2[0][0]']   \n",
      " icingOpLambda)                                                                                   \n",
      "                                                                                                  \n",
      " tf.__operators__.getitem_3 (Sl  ()                  0           ['tf.compat.v1.shape_3[0][0]']   \n",
      " icingOpLambda)                                                                                   \n",
      "                                                                                                  \n",
      " tf.nn.relu_2 (TFOpLambda)      (1, 1, 8, 64)        0           ['conv2_encoderBatchNorm[0][0]'] \n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "                                                                                                  \n",
      " tf.compat.v1.transpose (TFOpLa  (1, 1, 64, 48)      0           ['tf.nn.relu_5[0][0]']           \n",
      " mbda)                                                                                            \n",
      "                                                                                                  \n",
      " tf.__operators__.getitem (Slic  ()                  0           ['tf.compat.v1.shape[0][0]']     \n",
      " ingOpLambda)                                                                                     \n",
      "                                                                                                  \n",
      " tf.__operators__.getitem_1 (Sl  ()                  0           ['tf.compat.v1.shape_1[0][0]']   \n",
      " icingOpLambda)                                                                                   \n",
      "                                                                                                  \n",
      " tf.math.multiply (TFOpLambda)  ()                   0           ['tf.__operators__.getitem_2[0][0\n",
      "                                                                 ]',                              \n",
      "                                                                  'tf.__operators__.getitem_3[0][0\n",
      "                                                                 ]']                              \n",
      "                                                                                                  \n",
      " tf.compat.v1.pad_3 (TFOpLambda  (1, 1, 10, 64)      0           ['tf.nn.relu_2[0][0]']           \n",
      " )                                                                                                \n",
      "                                                                                                  \n",
      " tf.reshape (TFOpLambda)        (1, 1, 3072)         0           ['tf.compat.v1.transpose[0][0]', \n",
      "                                                                  'tf.__operators__.getitem[0][0]'\n",
      "                                                                 , 'tf.__operators__.getitem_1[0][\n",
      "                                                                 0]',                             \n",
      "                                                                  'tf.math.multiply[0][0]']       \n",
      "                                                                                                  \n",
      " conv3_encoder (Conv2D)         (1, 1, 8, 64)        192         ['tf.compat.v1.pad_3[0][0]']     \n",
      "                                                                                                  \n",
      " tf.split (TFOpLambda)          [(1, 1, 384),        0           ['tf.reshape[0][0]']             \n",
      "                                 (1, 1, 384),                                                     \n",
      "                                 (1, 1, 384),                                                     \n",
      "                                 (1, 1, 384),                                                     \n",
      "                                 (1, 1, 384),                                                     \n",
      "                                 (1, 1, 384),                                                     \n",
      "                                 (1, 1, 384),                                                     \n",
      "                                 (1, 1, 384)]                                                     \n",
      "                                                                                                  \n",
      " conv3_encoder_1x1 (Conv2D)     (1, 1, 8, 64)        4096        ['conv3_encoder[0][0]']          \n",
      "                                                                                                  \n",
      " GFC_encoder_0 (Dense)          (1, 1, 64)           24640       ['tf.split[0][0]']               \n",
      "                                                                                                  \n",
      " GFC_encoder_1 (Dense)          (1, 1, 64)           24640       ['tf.split[0][1]']               \n",
      "                                                                                                  \n",
      " GFC_encoder_2 (Dense)          (1, 1, 64)           24640       ['tf.split[0][2]']               \n",
      "                                                                                                  \n",
      " GFC_encoder_3 (Dense)          (1, 1, 64)           24640       ['tf.split[0][3]']               \n",
      "                                                                                                  \n",
      " GFC_encoder_4 (Dense)          (1, 1, 64)           24640       ['tf.split[0][4]']               \n",
      "                                                                                                  \n",
      " GFC_encoder_5 (Dense)          (1, 1, 64)           24640       ['tf.split[0][5]']               \n",
      "                                                                                                  \n",
      " GFC_encoder_6 (Dense)          (1, 1, 64)           24640       ['tf.split[0][6]']               \n",
      "                                                                                                  \n",
      " GFC_encoder_7 (Dense)          (1, 1, 64)           24640       ['tf.split[0][7]']               \n",
      "                                                                                                  \n",
      " conv3_encoderBatchNorm (BatchN  (1, 1, 8, 64)       256         ['conv3_encoder_1x1[0][0]']      \n",
      " ormalization)                                                                                    \n",
      "                                                                                                  \n",
      " tf.stack (TFOpLambda)          (1, 1, 64, 8)        0           ['GFC_encoder_0[0][0]',          \n",
      "                                                                  'GFC_encoder_1[0][0]',          \n",
      "                                                                  'GFC_encoder_2[0][0]',          \n",
      "                                                                  'GFC_encoder_3[0][0]',          \n",
      "                                                                  'GFC_encoder_4[0][0]',          \n",
      "                                                                  'GFC_encoder_5[0][0]',          \n",
      "                                                                  'GFC_encoder_6[0][0]',          \n",
      "                                                                  'GFC_encoder_7[0][0]']          \n",
      "                                                                                                  \n",
      " tf.nn.relu_3 (TFOpLambda)      (1, 1, 8, 64)        0           ['conv3_encoderBatchNorm[0][0]'] \n",
      "                                                                                                  \n",
      " tf.compat.v1.transpose_1 (TFOp  (1, 1, 8, 64)       0           ['tf.stack[0][0]']               \n",
      " Lambda)                                                                                          \n",
      "                                                                                                  \n",
      " tf.compat.v1.shape_10 (TFOpLam  (4,)                0           ['tf.nn.relu_3[0][0]']           \n",
      " bda)                                                                                             \n",
      "                                                                                                  \n",
      " tf.compat.v1.shape_11 (TFOpLam  (4,)                0           ['tf.nn.relu_3[0][0]']           \n",
      " bda)                                                                                             \n",
      "                                                                                                  \n",
      " tf.compat.v1.shape_7 (TFOpLamb  (4,)                0           ['tf.compat.v1.transpose_1[0][0]'\n",
      " da)                                                             ]                                \n",
      "                                                                                                  \n",
      " tf.compat.v1.shape_6 (TFOpLamb  (4,)                0           ['tf.compat.v1.transpose_1[0][0]'\n",
      " da)                                                             ]                                \n",
      "                                                                                                  \n",
      " tf.compat.v1.shape_8 (TFOpLamb  (4,)                0           ['tf.nn.relu_3[0][0]']           \n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      " da)                                                                                              \n",
      "                                                                                                  \n",
      " tf.compat.v1.shape_9 (TFOpLamb  (4,)                0           ['tf.nn.relu_3[0][0]']           \n",
      " da)                                                                                              \n",
      "                                                                                                  \n",
      " tf.__operators__.getitem_10 (S  ()                  0           ['tf.compat.v1.shape_10[0][0]']  \n",
      " licingOpLambda)                                                                                  \n",
      "                                                                                                  \n",
      " tf.__operators__.getitem_11 (S  ()                  0           ['tf.compat.v1.shape_11[0][0]']  \n",
      " licingOpLambda)                                                                                  \n",
      "                                                                                                  \n",
      " tf.compat.v1.shape_4 (TFOpLamb  (4,)                0           ['tf.compat.v1.transpose_1[0][0]'\n",
      " da)                                                             ]                                \n",
      "                                                                                                  \n",
      " tf.compat.v1.shape_5 (TFOpLamb  (4,)                0           ['tf.compat.v1.transpose_1[0][0]'\n",
      " da)                                                             ]                                \n",
      "                                                                                                  \n",
      " tf.__operators__.getitem_7 (Sl  ()                  0           ['tf.compat.v1.shape_7[0][0]']   \n",
      " icingOpLambda)                                                                                   \n",
      "                                                                                                  \n",
      " tf.__operators__.getitem_6 (Sl  ()                  0           ['tf.compat.v1.shape_6[0][0]']   \n",
      " icingOpLambda)                                                                                   \n",
      "                                                                                                  \n",
      " tf.compat.v1.transpose_2 (TFOp  (1, 1, 64, 8)       0           ['tf.nn.relu_3[0][0]']           \n",
      " Lambda)                                                                                          \n",
      "                                                                                                  \n",
      " tf.__operators__.getitem_8 (Sl  ()                  0           ['tf.compat.v1.shape_8[0][0]']   \n",
      " icingOpLambda)                                                                                   \n",
      "                                                                                                  \n",
      " tf.__operators__.getitem_9 (Sl  ()                  0           ['tf.compat.v1.shape_9[0][0]']   \n",
      " icingOpLambda)                                                                                   \n",
      "                                                                                                  \n",
      " tf.math.multiply_2 (TFOpLambda  ()                  0           ['tf.__operators__.getitem_10[0][\n",
      " )                                                               0]',                             \n",
      "                                                                  'tf.__operators__.getitem_11[0][\n",
      "                                                                 0]']                             \n",
      "                                                                                                  \n",
      " tf.__operators__.getitem_4 (Sl  ()                  0           ['tf.compat.v1.shape_4[0][0]']   \n",
      " icingOpLambda)                                                                                   \n",
      "                                                                                                  \n",
      " tf.__operators__.getitem_5 (Sl  ()                  0           ['tf.compat.v1.shape_5[0][0]']   \n",
      " icingOpLambda)                                                                                   \n",
      "                                                                                                  \n",
      " tf.math.multiply_1 (TFOpLambda  ()                  0           ['tf.__operators__.getitem_7[0][0\n",
      " )                                                               ]',                              \n",
      "                                                                  'tf.__operators__.getitem_6[0][0\n",
      "                                                                 ]']                              \n",
      "                                                                                                  \n",
      " tf.reshape_2 (TFOpLambda)      (1, 1, 512)          0           ['tf.compat.v1.transpose_2[0][0]'\n",
      "                                                                 , 'tf.__operators__.getitem_8[0][\n",
      "                                                                 0]',                             \n",
      "                                                                  'tf.__operators__.getitem_9[0][0\n",
      "                                                                 ]',                              \n",
      "                                                                  'tf.math.multiply_2[0][0]']     \n",
      "                                                                                                  \n",
      " tf.reshape_1 (TFOpLambda)      (1, 1, 512)          0           ['tf.compat.v1.transpose_1[0][0]'\n",
      "                                                                 , 'tf.__operators__.getitem_4[0][\n",
      "                                                                 0]',                             \n",
      "                                                                  'tf.__operators__.getitem_5[0][0\n",
      "                                                                 ]',                              \n",
      "                                                                  'tf.math.multiply_1[0][0]']     \n",
      "                                                                                                  \n",
      " concatenate (Concatenate)      (1, 1, 1024)         0           ['tf.reshape_2[0][0]',           \n",
      "                                                                  'tf.reshape_1[0][0]']           \n",
      "                                                                                                  \n",
      " state0 (InputLayer)            [(1, 256)]           0           []                               \n",
      "                                                                                                  \n",
      " GGRU00 (GRU)                   [(1, 1, 256),        984576      ['concatenate[0][0]',            \n",
      "                                 (1, 256)]                        'state0[0][0]']                 \n",
      "                                                                                                  \n",
      " tf.nn.relu_6 (TFOpLambda)      (1, 1, 256)          0           ['GGRU00[0][0]']                 \n",
      "                                                                                                  \n",
      " state10 (InputLayer)           [(1, 256)]           0           []                               \n",
      "                                                                                                  \n",
      " GGRU010 (GRU)                  [(1, 1, 256),        394752      ['tf.nn.relu_6[0][0]',           \n",
      "                                 (1, 256)]                        'state10[0][0]']                \n",
      "                                                                                                  \n",
      " tf.nn.relu_7 (TFOpLambda)      (1, 1, 256)          0           ['GGRU010[0][0]']                \n",
      "                                                                                                  \n",
      " state11 (InputLayer)           [(1, 256)]           0           []                               \n",
      "                                                                                                  \n",
      " GGRU011 (GRU)                  [(1, 1, 256),        394752      ['tf.nn.relu_7[0][0]',           \n",
      "                                 (1, 256)]                        'state11[0][0]']                \n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "                                                                                                  \n",
      " tf.nn.relu_8 (TFOpLambda)      (1, 1, 256)          0           ['GGRU011[0][0]']                \n",
      "                                                                                                  \n",
      " tf.__operators__.add (TFOpLamb  (1, 1, 256)         0           ['tf.nn.relu_7[0][0]',           \n",
      " da)                                                              'tf.nn.relu_8[0][0]']           \n",
      "                                                                                                  \n",
      " tf.split_1 (TFOpLambda)        [(1, 1, 32),         0           ['tf.__operators__.add[0][0]']   \n",
      "                                 (1, 1, 32),                                                      \n",
      "                                 (1, 1, 32),                                                      \n",
      "                                 (1, 1, 32),                                                      \n",
      "                                 (1, 1, 32),                                                      \n",
      "                                 (1, 1, 32),                                                      \n",
      "                                 (1, 1, 32),                                                      \n",
      "                                 (1, 1, 32)]                                                      \n",
      "                                                                                                  \n",
      " GFC_decoder_0 (Dense)          (1, 1, 64)           2112        ['tf.split_1[0][0]']             \n",
      "                                                                                                  \n",
      " GFC_decoder_1 (Dense)          (1, 1, 64)           2112        ['tf.split_1[0][1]']             \n",
      "                                                                                                  \n",
      " GFC_decoder_2 (Dense)          (1, 1, 64)           2112        ['tf.split_1[0][2]']             \n",
      "                                                                                                  \n",
      " GFC_decoder_3 (Dense)          (1, 1, 64)           2112        ['tf.split_1[0][3]']             \n",
      "                                                                                                  \n",
      " GFC_decoder_4 (Dense)          (1, 1, 64)           2112        ['tf.split_1[0][4]']             \n",
      "                                                                                                  \n",
      " GFC_decoder_5 (Dense)          (1, 1, 64)           2112        ['tf.split_1[0][5]']             \n",
      "                                                                                                  \n",
      " GFC_decoder_6 (Dense)          (1, 1, 64)           2112        ['tf.split_1[0][6]']             \n",
      "                                                                                                  \n",
      " GFC_decoder_7 (Dense)          (1, 1, 64)           2112        ['tf.split_1[0][7]']             \n",
      "                                                                                                  \n",
      " tf.stack_1 (TFOpLambda)        (1, 1, 64, 8)        0           ['GFC_decoder_0[0][0]',          \n",
      "                                                                  'GFC_decoder_1[0][0]',          \n",
      "                                                                  'GFC_decoder_2[0][0]',          \n",
      "                                                                  'GFC_decoder_3[0][0]',          \n",
      "                                                                  'GFC_decoder_4[0][0]',          \n",
      "                                                                  'GFC_decoder_5[0][0]',          \n",
      "                                                                  'GFC_decoder_6[0][0]',          \n",
      "                                                                  'GFC_decoder_7[0][0]']          \n",
      "                                                                                                  \n",
      " tf.compat.v1.transpose_3 (TFOp  (1, 1, 8, 64)       0           ['tf.stack_1[0][0]']             \n",
      " Lambda)                                                                                          \n",
      "                                                                                                  \n",
      " tf.compat.v1.shape_15 (TFOpLam  (4,)                0           ['tf.compat.v1.transpose_3[0][0]'\n",
      " bda)                                                            ]                                \n",
      "                                                                                                  \n",
      " tf.compat.v1.shape_14 (TFOpLam  (4,)                0           ['tf.compat.v1.transpose_3[0][0]'\n",
      " bda)                                                            ]                                \n",
      "                                                                                                  \n",
      " tf.compat.v1.shape_12 (TFOpLam  (4,)                0           ['tf.compat.v1.transpose_3[0][0]'\n",
      " bda)                                                            ]                                \n",
      "                                                                                                  \n",
      " tf.compat.v1.shape_13 (TFOpLam  (4,)                0           ['tf.compat.v1.transpose_3[0][0]'\n",
      " bda)                                                            ]                                \n",
      "                                                                                                  \n",
      " tf.__operators__.getitem_15 (S  ()                  0           ['tf.compat.v1.shape_15[0][0]']  \n",
      " licingOpLambda)                                                                                  \n",
      "                                                                                                  \n",
      " tf.__operators__.getitem_14 (S  ()                  0           ['tf.compat.v1.shape_14[0][0]']  \n",
      " licingOpLambda)                                                                                  \n",
      "                                                                                                  \n",
      " tf.__operators__.getitem_12 (S  ()                  0           ['tf.compat.v1.shape_12[0][0]']  \n",
      " licingOpLambda)                                                                                  \n",
      "                                                                                                  \n",
      " tf.__operators__.getitem_13 (S  ()                  0           ['tf.compat.v1.shape_13[0][0]']  \n",
      " licingOpLambda)                                                                                  \n",
      "                                                                                                  \n",
      " tf.math.multiply_3 (TFOpLambda  ()                  0           ['tf.__operators__.getitem_15[0][\n",
      " )                                                               0]',                             \n",
      "                                                                  'tf.__operators__.getitem_14[0][\n",
      "                                                                 0]']                             \n",
      "                                                                                                  \n",
      " convp3 (Conv2D)                (1, 1, 8, 64)        64          ['tf.nn.relu_3[0][0]']           \n",
      "                                                                                                  \n",
      " tf.reshape_3 (TFOpLambda)      (1, 1, 512)          0           ['tf.compat.v1.transpose_3[0][0]'\n",
      "                                                                 , 'tf.__operators__.getitem_12[0]\n",
      "                                                                 [0]',                            \n",
      "                                                                  'tf.__operators__.getitem_13[0][\n",
      "                                                                 0]',                             \n",
      "                                                                  'tf.math.multiply_3[0][0]']     \n",
      "                                                                                                  \n",
      " tf.compat.v1.shape_16 (TFOpLam  (4,)                0           ['tf.nn.relu_3[0][0]']           \n",
      " bda)                                                                                             \n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "                                                                                                  \n",
      " tf.compat.v1.shape_17 (TFOpLam  (4,)                0           ['tf.nn.relu_3[0][0]']           \n",
      " bda)                                                                                             \n",
      "                                                                                                  \n",
      " tf.compat.v1.shape_18 (TFOpLam  (4,)                0           ['tf.nn.relu_3[0][0]']           \n",
      " bda)                                                                                             \n",
      "                                                                                                  \n",
      " convp3BatchNorm (BatchNormaliz  (1, 1, 8, 64)       256         ['convp3[0][0]']                 \n",
      " ation)                                                                                           \n",
      "                                                                                                  \n",
      " tf.nn.relu_9 (TFOpLambda)      (1, 1, 512)          0           ['tf.reshape_3[0][0]']           \n",
      "                                                                                                  \n",
      " tf.__operators__.getitem_16 (S  ()                  0           ['tf.compat.v1.shape_16[0][0]']  \n",
      " licingOpLambda)                                                                                  \n",
      "                                                                                                  \n",
      " tf.__operators__.getitem_17 (S  ()                  0           ['tf.compat.v1.shape_17[0][0]']  \n",
      " licingOpLambda)                                                                                  \n",
      "                                                                                                  \n",
      " tf.__operators__.getitem_18 (S  ()                  0           ['tf.compat.v1.shape_18[0][0]']  \n",
      " licingOpLambda)                                                                                  \n",
      "                                                                                                  \n",
      " tf.nn.relu_10 (TFOpLambda)     (1, 1, 8, 64)        0           ['convp3BatchNorm[0][0]']        \n",
      "                                                                                                  \n",
      " tf.reshape_4 (TFOpLambda)      (1, 1, 8, 64)        0           ['tf.nn.relu_9[0][0]',           \n",
      "                                                                  'tf.__operators__.getitem_16[0][\n",
      "                                                                 0]',                             \n",
      "                                                                  'tf.__operators__.getitem_17[0][\n",
      "                                                                 0]',                             \n",
      "                                                                  'tf.__operators__.getitem_18[0][\n",
      "                                                                 0]']                             \n",
      "                                                                                                  \n",
      " tf.__operators__.add_1 (TFOpLa  (1, 1, 8, 64)       0           ['tf.nn.relu_10[0][0]',          \n",
      " mbda)                                                            'tf.reshape_4[0][0]']           \n",
      "                                                                                                  \n",
      " tf.compat.v1.pad_6 (TFOpLambda  (1, 1, 10, 64)      0           ['tf.__operators__.add_1[0][0]'] \n",
      " )                                                                                                \n",
      "                                                                                                  \n",
      " convt3 (Conv2D)                (1, 1, 8, 64)        192         ['tf.compat.v1.pad_6[0][0]']     \n",
      "                                                                                                  \n",
      " convp2 (Conv2D)                (1, 1, 8, 64)        64          ['tf.nn.relu_2[0][0]']           \n",
      "                                                                                                  \n",
      " convt3_1x1 (Conv2D)            (1, 1, 8, 64)        4096        ['convt3[0][0]']                 \n",
      "                                                                                                  \n",
      " convp2BatchNorm (BatchNormaliz  (1, 1, 8, 64)       256         ['convp2[0][0]']                 \n",
      " ation)                                                                                           \n",
      "                                                                                                  \n",
      " convt3BatchNorm (BatchNormaliz  (1, 1, 8, 64)       256         ['convt3_1x1[0][0]']             \n",
      " ation)                                                                                           \n",
      "                                                                                                  \n",
      " tf.nn.relu_12 (TFOpLambda)     (1, 1, 8, 64)        0           ['convp2BatchNorm[0][0]']        \n",
      "                                                                                                  \n",
      " tf.nn.relu_11 (TFOpLambda)     (1, 1, 8, 64)        0           ['convt3BatchNorm[0][0]']        \n",
      "                                                                                                  \n",
      " tf.__operators__.add_2 (TFOpLa  (1, 1, 8, 64)       0           ['tf.nn.relu_12[0][0]',          \n",
      " mbda)                                                            'tf.nn.relu_11[0][0]']          \n",
      "                                                                                                  \n",
      " state20 (InputLayer)           [(1, 256)]           0           []                               \n",
      "                                                                                                  \n",
      " tf.split_2 (TFOpLambda)        [(1, 1, 8, 1),       0           ['tf.__operators__.add_2[0][0]'] \n",
      "                                 (1, 1, 8, 1),                                                    \n",
      "                                 (1, 1, 8, 1),                                                    \n",
      "                                 (1, 1, 8, 1),                                                    \n",
      "                                 (1, 1, 8, 1),                                                    \n",
      "                                 (1, 1, 8, 1),                                                    \n",
      "                                 (1, 1, 8, 1),                                                    \n",
      "                                 (1, 1, 8, 1),                                                    \n",
      "                                 (1, 1, 8, 1),                                                    \n",
      "                                 (1, 1, 8, 1),                                                    \n",
      "                                 (1, 1, 8, 1),                                                    \n",
      "                                 (1, 1, 8, 1),                                                    \n",
      "                                 (1, 1, 8, 1),                                                    \n",
      "                                 (1, 1, 8, 1),                                                    \n",
      "                                 (1, 1, 8, 1),                                                    \n",
      "                                 (1, 1, 8, 1),                                                    \n",
      "                                 (1, 1, 8, 1),                                                    \n",
      "                                 (1, 1, 8, 1),                                                    \n",
      "                                 (1, 1, 8, 1),                                                    \n",
      "                                 (1, 1, 8, 1),                                                    \n",
      "                                 (1, 1, 8, 1),                                                    \n",
      "                                 (1, 1, 8, 1),                                                    \n",
      "                                 (1, 1, 8, 1),                                                    \n",
      "                                 (1, 1, 8, 1),                                                    \n",
      "                                 (1, 1, 8, 1),                                                    \n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "                                 (1, 1, 8, 1),                                                    \n",
      "                                 (1, 1, 8, 1),                                                    \n",
      "                                 (1, 1, 8, 1),                                                    \n",
      "                                 (1, 1, 8, 1),                                                    \n",
      "                                 (1, 1, 8, 1),                                                    \n",
      "                                 (1, 1, 8, 1),                                                    \n",
      "                                 (1, 1, 8, 1),                                                    \n",
      "                                 (1, 1, 8, 1),                                                    \n",
      "                                 (1, 1, 8, 1),                                                    \n",
      "                                 (1, 1, 8, 1),                                                    \n",
      "                                 (1, 1, 8, 1),                                                    \n",
      "                                 (1, 1, 8, 1),                                                    \n",
      "                                 (1, 1, 8, 1),                                                    \n",
      "                                 (1, 1, 8, 1),                                                    \n",
      "                                 (1, 1, 8, 1),                                                    \n",
      "                                 (1, 1, 8, 1),                                                    \n",
      "                                 (1, 1, 8, 1),                                                    \n",
      "                                 (1, 1, 8, 1),                                                    \n",
      "                                 (1, 1, 8, 1),                                                    \n",
      "                                 (1, 1, 8, 1),                                                    \n",
      "                                 (1, 1, 8, 1),                                                    \n",
      "                                 (1, 1, 8, 1),                                                    \n",
      "                                 (1, 1, 8, 1),                                                    \n",
      "                                 (1, 1, 8, 1),                                                    \n",
      "                                 (1, 1, 8, 1),                                                    \n",
      "                                 (1, 1, 8, 1),                                                    \n",
      "                                 (1, 1, 8, 1),                                                    \n",
      "                                 (1, 1, 8, 1),                                                    \n",
      "                                 (1, 1, 8, 1),                                                    \n",
      "                                 (1, 1, 8, 1),                                                    \n",
      "                                 (1, 1, 8, 1),                                                    \n",
      "                                 (1, 1, 8, 1),                                                    \n",
      "                                 (1, 1, 8, 1),                                                    \n",
      "                                 (1, 1, 8, 1),                                                    \n",
      "                                 (1, 1, 8, 1),                                                    \n",
      "                                 (1, 1, 8, 1),                                                    \n",
      "                                 (1, 1, 8, 1),                                                    \n",
      "                                 (1, 1, 8, 1),                                                    \n",
      "                                 (1, 1, 8, 1)]                                                    \n",
      "                                                                                                  \n",
      " GGRU10 (GRU)                   [(1, 1, 256),        394752      ['tf.nn.relu_6[0][0]',           \n",
      "                                 (1, 256)]                        'state20[0][0]']                \n",
      "                                                                                                  \n",
      " conv2d_transpose (Conv2DTransp  (1, 1, 16, 1)       3           ['tf.split_2[0][0]']             \n",
      " ose)                                                                                             \n",
      "                                                                                                  \n",
      " conv2d_transpose_1 (Conv2DTran  (1, 1, 16, 1)       3           ['tf.split_2[0][1]']             \n",
      " spose)                                                                                           \n",
      "                                                                                                  \n",
      " conv2d_transpose_2 (Conv2DTran  (1, 1, 16, 1)       3           ['tf.split_2[0][2]']             \n",
      " spose)                                                                                           \n",
      "                                                                                                  \n",
      " conv2d_transpose_3 (Conv2DTran  (1, 1, 16, 1)       3           ['tf.split_2[0][3]']             \n",
      " spose)                                                                                           \n",
      "                                                                                                  \n",
      " conv2d_transpose_4 (Conv2DTran  (1, 1, 16, 1)       3           ['tf.split_2[0][4]']             \n",
      " spose)                                                                                           \n",
      "                                                                                                  \n",
      " conv2d_transpose_5 (Conv2DTran  (1, 1, 16, 1)       3           ['tf.split_2[0][5]']             \n",
      " spose)                                                                                           \n",
      "                                                                                                  \n",
      " conv2d_transpose_6 (Conv2DTran  (1, 1, 16, 1)       3           ['tf.split_2[0][6]']             \n",
      " spose)                                                                                           \n",
      "                                                                                                  \n",
      " conv2d_transpose_7 (Conv2DTran  (1, 1, 16, 1)       3           ['tf.split_2[0][7]']             \n",
      " spose)                                                                                           \n",
      "                                                                                                  \n",
      " conv2d_transpose_8 (Conv2DTran  (1, 1, 16, 1)       3           ['tf.split_2[0][8]']             \n",
      " spose)                                                                                           \n",
      "                                                                                                  \n",
      " conv2d_transpose_9 (Conv2DTran  (1, 1, 16, 1)       3           ['tf.split_2[0][9]']             \n",
      " spose)                                                                                           \n",
      "                                                                                                  \n",
      " conv2d_transpose_10 (Conv2DTra  (1, 1, 16, 1)       3           ['tf.split_2[0][10]']            \n",
      " nspose)                                                                                          \n",
      "                                                                                                  \n",
      " conv2d_transpose_11 (Conv2DTra  (1, 1, 16, 1)       3           ['tf.split_2[0][11]']            \n",
      " nspose)                                                                                          \n",
      "                                                                                                  \n",
      " conv2d_transpose_12 (Conv2DTra  (1, 1, 16, 1)       3           ['tf.split_2[0][12]']            \n",
      " nspose)                                                                                          \n",
      "                                                                                                  \n",
      " conv2d_transpose_13 (Conv2DTra  (1, 1, 16, 1)       3           ['tf.split_2[0][13]']            \n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      " nspose)                                                                                          \n",
      "                                                                                                  \n",
      " conv2d_transpose_14 (Conv2DTra  (1, 1, 16, 1)       3           ['tf.split_2[0][14]']            \n",
      " nspose)                                                                                          \n",
      "                                                                                                  \n",
      " conv2d_transpose_15 (Conv2DTra  (1, 1, 16, 1)       3           ['tf.split_2[0][15]']            \n",
      " nspose)                                                                                          \n",
      "                                                                                                  \n",
      " conv2d_transpose_16 (Conv2DTra  (1, 1, 16, 1)       3           ['tf.split_2[0][16]']            \n",
      " nspose)                                                                                          \n",
      "                                                                                                  \n",
      " conv2d_transpose_17 (Conv2DTra  (1, 1, 16, 1)       3           ['tf.split_2[0][17]']            \n",
      " nspose)                                                                                          \n",
      "                                                                                                  \n",
      " conv2d_transpose_18 (Conv2DTra  (1, 1, 16, 1)       3           ['tf.split_2[0][18]']            \n",
      " nspose)                                                                                          \n",
      "                                                                                                  \n",
      " conv2d_transpose_19 (Conv2DTra  (1, 1, 16, 1)       3           ['tf.split_2[0][19]']            \n",
      " nspose)                                                                                          \n",
      "                                                                                                  \n",
      " conv2d_transpose_20 (Conv2DTra  (1, 1, 16, 1)       3           ['tf.split_2[0][20]']            \n",
      " nspose)                                                                                          \n",
      "                                                                                                  \n",
      " conv2d_transpose_21 (Conv2DTra  (1, 1, 16, 1)       3           ['tf.split_2[0][21]']            \n",
      " nspose)                                                                                          \n",
      "                                                                                                  \n",
      " conv2d_transpose_22 (Conv2DTra  (1, 1, 16, 1)       3           ['tf.split_2[0][22]']            \n",
      " nspose)                                                                                          \n",
      "                                                                                                  \n",
      " conv2d_transpose_23 (Conv2DTra  (1, 1, 16, 1)       3           ['tf.split_2[0][23]']            \n",
      " nspose)                                                                                          \n",
      "                                                                                                  \n",
      " conv2d_transpose_24 (Conv2DTra  (1, 1, 16, 1)       3           ['tf.split_2[0][24]']            \n",
      " nspose)                                                                                          \n",
      "                                                                                                  \n",
      " conv2d_transpose_25 (Conv2DTra  (1, 1, 16, 1)       3           ['tf.split_2[0][25]']            \n",
      " nspose)                                                                                          \n",
      "                                                                                                  \n",
      " conv2d_transpose_26 (Conv2DTra  (1, 1, 16, 1)       3           ['tf.split_2[0][26]']            \n",
      " nspose)                                                                                          \n",
      "                                                                                                  \n",
      " conv2d_transpose_27 (Conv2DTra  (1, 1, 16, 1)       3           ['tf.split_2[0][27]']            \n",
      " nspose)                                                                                          \n",
      "                                                                                                  \n",
      " conv2d_transpose_28 (Conv2DTra  (1, 1, 16, 1)       3           ['tf.split_2[0][28]']            \n",
      " nspose)                                                                                          \n",
      "                                                                                                  \n",
      " conv2d_transpose_29 (Conv2DTra  (1, 1, 16, 1)       3           ['tf.split_2[0][29]']            \n",
      " nspose)                                                                                          \n",
      "                                                                                                  \n",
      " conv2d_transpose_30 (Conv2DTra  (1, 1, 16, 1)       3           ['tf.split_2[0][30]']            \n",
      " nspose)                                                                                          \n",
      "                                                                                                  \n",
      " conv2d_transpose_31 (Conv2DTra  (1, 1, 16, 1)       3           ['tf.split_2[0][31]']            \n",
      " nspose)                                                                                          \n",
      "                                                                                                  \n",
      " conv2d_transpose_32 (Conv2DTra  (1, 1, 16, 1)       3           ['tf.split_2[0][32]']            \n",
      " nspose)                                                                                          \n",
      "                                                                                                  \n",
      " conv2d_transpose_33 (Conv2DTra  (1, 1, 16, 1)       3           ['tf.split_2[0][33]']            \n",
      " nspose)                                                                                          \n",
      "                                                                                                  \n",
      " conv2d_transpose_34 (Conv2DTra  (1, 1, 16, 1)       3           ['tf.split_2[0][34]']            \n",
      " nspose)                                                                                          \n",
      "                                                                                                  \n",
      " conv2d_transpose_35 (Conv2DTra  (1, 1, 16, 1)       3           ['tf.split_2[0][35]']            \n",
      " nspose)                                                                                          \n",
      "                                                                                                  \n",
      " conv2d_transpose_36 (Conv2DTra  (1, 1, 16, 1)       3           ['tf.split_2[0][36]']            \n",
      " nspose)                                                                                          \n",
      "                                                                                                  \n",
      " conv2d_transpose_37 (Conv2DTra  (1, 1, 16, 1)       3           ['tf.split_2[0][37]']            \n",
      " nspose)                                                                                          \n",
      "                                                                                                  \n",
      " conv2d_transpose_38 (Conv2DTra  (1, 1, 16, 1)       3           ['tf.split_2[0][38]']            \n",
      " nspose)                                                                                          \n",
      "                                                                                                  \n",
      " conv2d_transpose_39 (Conv2DTra  (1, 1, 16, 1)       3           ['tf.split_2[0][39]']            \n",
      " nspose)                                                                                          \n",
      "                                                                                                  \n",
      " conv2d_transpose_40 (Conv2DTra  (1, 1, 16, 1)       3           ['tf.split_2[0][40]']            \n",
      " nspose)                                                                                          \n",
      "                                                                                                  \n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      " conv2d_transpose_41 (Conv2DTra  (1, 1, 16, 1)       3           ['tf.split_2[0][41]']            \n",
      " nspose)                                                                                          \n",
      "                                                                                                  \n",
      " conv2d_transpose_42 (Conv2DTra  (1, 1, 16, 1)       3           ['tf.split_2[0][42]']            \n",
      " nspose)                                                                                          \n",
      "                                                                                                  \n",
      " conv2d_transpose_43 (Conv2DTra  (1, 1, 16, 1)       3           ['tf.split_2[0][43]']            \n",
      " nspose)                                                                                          \n",
      "                                                                                                  \n",
      " conv2d_transpose_44 (Conv2DTra  (1, 1, 16, 1)       3           ['tf.split_2[0][44]']            \n",
      " nspose)                                                                                          \n",
      "                                                                                                  \n",
      " conv2d_transpose_45 (Conv2DTra  (1, 1, 16, 1)       3           ['tf.split_2[0][45]']            \n",
      " nspose)                                                                                          \n",
      "                                                                                                  \n",
      " conv2d_transpose_46 (Conv2DTra  (1, 1, 16, 1)       3           ['tf.split_2[0][46]']            \n",
      " nspose)                                                                                          \n",
      "                                                                                                  \n",
      " conv2d_transpose_47 (Conv2DTra  (1, 1, 16, 1)       3           ['tf.split_2[0][47]']            \n",
      " nspose)                                                                                          \n",
      "                                                                                                  \n",
      " conv2d_transpose_48 (Conv2DTra  (1, 1, 16, 1)       3           ['tf.split_2[0][48]']            \n",
      " nspose)                                                                                          \n",
      "                                                                                                  \n",
      " conv2d_transpose_49 (Conv2DTra  (1, 1, 16, 1)       3           ['tf.split_2[0][49]']            \n",
      " nspose)                                                                                          \n",
      "                                                                                                  \n",
      " conv2d_transpose_50 (Conv2DTra  (1, 1, 16, 1)       3           ['tf.split_2[0][50]']            \n",
      " nspose)                                                                                          \n",
      "                                                                                                  \n",
      " conv2d_transpose_51 (Conv2DTra  (1, 1, 16, 1)       3           ['tf.split_2[0][51]']            \n",
      " nspose)                                                                                          \n",
      "                                                                                                  \n",
      " conv2d_transpose_52 (Conv2DTra  (1, 1, 16, 1)       3           ['tf.split_2[0][52]']            \n",
      " nspose)                                                                                          \n",
      "                                                                                                  \n",
      " conv2d_transpose_53 (Conv2DTra  (1, 1, 16, 1)       3           ['tf.split_2[0][53]']            \n",
      " nspose)                                                                                          \n",
      "                                                                                                  \n",
      " conv2d_transpose_54 (Conv2DTra  (1, 1, 16, 1)       3           ['tf.split_2[0][54]']            \n",
      " nspose)                                                                                          \n",
      "                                                                                                  \n",
      " conv2d_transpose_55 (Conv2DTra  (1, 1, 16, 1)       3           ['tf.split_2[0][55]']            \n",
      " nspose)                                                                                          \n",
      "                                                                                                  \n",
      " conv2d_transpose_56 (Conv2DTra  (1, 1, 16, 1)       3           ['tf.split_2[0][56]']            \n",
      " nspose)                                                                                          \n",
      "                                                                                                  \n",
      " conv2d_transpose_57 (Conv2DTra  (1, 1, 16, 1)       3           ['tf.split_2[0][57]']            \n",
      " nspose)                                                                                          \n",
      "                                                                                                  \n",
      " conv2d_transpose_58 (Conv2DTra  (1, 1, 16, 1)       3           ['tf.split_2[0][58]']            \n",
      " nspose)                                                                                          \n",
      "                                                                                                  \n",
      " conv2d_transpose_59 (Conv2DTra  (1, 1, 16, 1)       3           ['tf.split_2[0][59]']            \n",
      " nspose)                                                                                          \n",
      "                                                                                                  \n",
      " conv2d_transpose_60 (Conv2DTra  (1, 1, 16, 1)       3           ['tf.split_2[0][60]']            \n",
      " nspose)                                                                                          \n",
      "                                                                                                  \n",
      " conv2d_transpose_61 (Conv2DTra  (1, 1, 16, 1)       3           ['tf.split_2[0][61]']            \n",
      " nspose)                                                                                          \n",
      "                                                                                                  \n",
      " conv2d_transpose_62 (Conv2DTra  (1, 1, 16, 1)       3           ['tf.split_2[0][62]']            \n",
      " nspose)                                                                                          \n",
      "                                                                                                  \n",
      " conv2d_transpose_63 (Conv2DTra  (1, 1, 16, 1)       3           ['tf.split_2[0][63]']            \n",
      " nspose)                                                                                          \n",
      "                                                                                                  \n",
      " tf.nn.relu_17 (TFOpLambda)     (1, 1, 256)          0           ['GGRU10[0][0]']                 \n",
      "                                                                                                  \n",
      " state21 (InputLayer)           [(1, 256)]           0           []                               \n",
      "                                                                                                  \n",
      " tf.concat (TFOpLambda)         (1, 1, 16, 64)       0           ['conv2d_transpose[0][0]',       \n",
      "                                                                  'conv2d_transpose_1[0][0]',     \n",
      "                                                                  'conv2d_transpose_2[0][0]',     \n",
      "                                                                  'conv2d_transpose_3[0][0]',     \n",
      "                                                                  'conv2d_transpose_4[0][0]',     \n",
      "                                                                  'conv2d_transpose_5[0][0]',     \n",
      "                                                                  'conv2d_transpose_6[0][0]',     \n",
      "                                                                  'conv2d_transpose_7[0][0]',     \n",
      "                                                                  'conv2d_transpose_8[0][0]',     \n",
      "                                                                  'conv2d_transpose_9[0][0]',     \n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "                                                                  'conv2d_transpose_10[0][0]',    \n",
      "                                                                  'conv2d_transpose_11[0][0]',    \n",
      "                                                                  'conv2d_transpose_12[0][0]',    \n",
      "                                                                  'conv2d_transpose_13[0][0]',    \n",
      "                                                                  'conv2d_transpose_14[0][0]',    \n",
      "                                                                  'conv2d_transpose_15[0][0]',    \n",
      "                                                                  'conv2d_transpose_16[0][0]',    \n",
      "                                                                  'conv2d_transpose_17[0][0]',    \n",
      "                                                                  'conv2d_transpose_18[0][0]',    \n",
      "                                                                  'conv2d_transpose_19[0][0]',    \n",
      "                                                                  'conv2d_transpose_20[0][0]',    \n",
      "                                                                  'conv2d_transpose_21[0][0]',    \n",
      "                                                                  'conv2d_transpose_22[0][0]',    \n",
      "                                                                  'conv2d_transpose_23[0][0]',    \n",
      "                                                                  'conv2d_transpose_24[0][0]',    \n",
      "                                                                  'conv2d_transpose_25[0][0]',    \n",
      "                                                                  'conv2d_transpose_26[0][0]',    \n",
      "                                                                  'conv2d_transpose_27[0][0]',    \n",
      "                                                                  'conv2d_transpose_28[0][0]',    \n",
      "                                                                  'conv2d_transpose_29[0][0]',    \n",
      "                                                                  'conv2d_transpose_30[0][0]',    \n",
      "                                                                  'conv2d_transpose_31[0][0]',    \n",
      "                                                                  'conv2d_transpose_32[0][0]',    \n",
      "                                                                  'conv2d_transpose_33[0][0]',    \n",
      "                                                                  'conv2d_transpose_34[0][0]',    \n",
      "                                                                  'conv2d_transpose_35[0][0]',    \n",
      "                                                                  'conv2d_transpose_36[0][0]',    \n",
      "                                                                  'conv2d_transpose_37[0][0]',    \n",
      "                                                                  'conv2d_transpose_38[0][0]',    \n",
      "                                                                  'conv2d_transpose_39[0][0]',    \n",
      "                                                                  'conv2d_transpose_40[0][0]',    \n",
      "                                                                  'conv2d_transpose_41[0][0]',    \n",
      "                                                                  'conv2d_transpose_42[0][0]',    \n",
      "                                                                  'conv2d_transpose_43[0][0]',    \n",
      "                                                                  'conv2d_transpose_44[0][0]',    \n",
      "                                                                  'conv2d_transpose_45[0][0]',    \n",
      "                                                                  'conv2d_transpose_46[0][0]',    \n",
      "                                                                  'conv2d_transpose_47[0][0]',    \n",
      "                                                                  'conv2d_transpose_48[0][0]',    \n",
      "                                                                  'conv2d_transpose_49[0][0]',    \n",
      "                                                                  'conv2d_transpose_50[0][0]',    \n",
      "                                                                  'conv2d_transpose_51[0][0]',    \n",
      "                                                                  'conv2d_transpose_52[0][0]',    \n",
      "                                                                  'conv2d_transpose_53[0][0]',    \n",
      "                                                                  'conv2d_transpose_54[0][0]',    \n",
      "                                                                  'conv2d_transpose_55[0][0]',    \n",
      "                                                                  'conv2d_transpose_56[0][0]',    \n",
      "                                                                  'conv2d_transpose_57[0][0]',    \n",
      "                                                                  'conv2d_transpose_58[0][0]',    \n",
      "                                                                  'conv2d_transpose_59[0][0]',    \n",
      "                                                                  'conv2d_transpose_60[0][0]',    \n",
      "                                                                  'conv2d_transpose_61[0][0]',    \n",
      "                                                                  'conv2d_transpose_62[0][0]',    \n",
      "                                                                  'conv2d_transpose_63[0][0]']    \n",
      "                                                                                                  \n",
      " GGRU11 (GRU)                   [(1, 1, 256),        394752      ['tf.nn.relu_17[0][0]',          \n",
      "                                 (1, 256)]                        'state21[0][0]']                \n",
      "                                                                                                  \n",
      " convp1 (Conv2D)                (1, 1, 16, 64)       64          ['tf.nn.relu_1[0][0]']           \n",
      "                                                                                                  \n",
      " convt2_1x1 (Conv2D)            (1, 1, 16, 64)       4096        ['tf.concat[0][0]']              \n",
      "                                                                                                  \n",
      " tf.nn.relu_18 (TFOpLambda)     (1, 1, 256)          0           ['GGRU11[0][0]']                 \n",
      "                                                                                                  \n",
      " convp1BatchNorm (BatchNormaliz  (1, 1, 16, 64)      256         ['convp1[0][0]']                 \n",
      " ation)                                                                                           \n",
      "                                                                                                  \n",
      " convt2BatchNorm (BatchNormaliz  (1, 1, 16, 64)      256         ['convt2_1x1[0][0]']             \n",
      " ation)                                                                                           \n",
      "                                                                                                  \n",
      " tf.__operators__.add_5 (TFOpLa  (1, 1, 256)         0           ['tf.nn.relu_17[0][0]',          \n",
      " mbda)                                                            'tf.nn.relu_18[0][0]']          \n",
      "                                                                                                  \n",
      " tf.nn.relu_14 (TFOpLambda)     (1, 1, 16, 64)       0           ['convp1BatchNorm[0][0]']        \n",
      "                                                                                                  \n",
      " tf.nn.relu_13 (TFOpLambda)     (1, 1, 16, 64)       0           ['convt2BatchNorm[0][0]']        \n",
      "                                                                                                  \n",
      " tf.split_4 (TFOpLambda)        [(1, 1, 32),         0           ['tf.__operators__.add_5[0][0]'] \n",
      "                                 (1, 1, 32),                                                      \n",
      "                                 (1, 1, 32),                                                      \n",
      "                                 (1, 1, 32),                                                      \n",
      "                                 (1, 1, 32),                                                      \n",
      "                                 (1, 1, 32),                                                      \n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "                                 (1, 1, 32),                                                      \n",
      "                                 (1, 1, 32)]                                                      \n",
      "                                                                                                  \n",
      " tf.__operators__.add_3 (TFOpLa  (1, 1, 16, 64)      0           ['tf.nn.relu_14[0][0]',          \n",
      " mbda)                                                            'tf.nn.relu_13[0][0]']          \n",
      "                                                                                                  \n",
      " GFC_c_0 (Dense)                (1, 1, 120)          3960        ['tf.split_4[0][0]']             \n",
      "                                                                                                  \n",
      " GFC_c_1 (Dense)                (1, 1, 120)          3960        ['tf.split_4[0][1]']             \n",
      "                                                                                                  \n",
      " GFC_c_2 (Dense)                (1, 1, 120)          3960        ['tf.split_4[0][2]']             \n",
      "                                                                                                  \n",
      " GFC_c_3 (Dense)                (1, 1, 120)          3960        ['tf.split_4[0][3]']             \n",
      "                                                                                                  \n",
      " GFC_c_4 (Dense)                (1, 1, 120)          3960        ['tf.split_4[0][4]']             \n",
      "                                                                                                  \n",
      " GFC_c_5 (Dense)                (1, 1, 120)          3960        ['tf.split_4[0][5]']             \n",
      "                                                                                                  \n",
      " GFC_c_6 (Dense)                (1, 1, 120)          3960        ['tf.split_4[0][6]']             \n",
      "                                                                                                  \n",
      " GFC_c_7 (Dense)                (1, 1, 120)          3960        ['tf.split_4[0][7]']             \n",
      "                                                                                                  \n",
      " tf.split_3 (TFOpLambda)        [(1, 1, 16, 1),      0           ['tf.__operators__.add_3[0][0]'] \n",
      "                                 (1, 1, 16, 1),                                                   \n",
      "                                 (1, 1, 16, 1),                                                   \n",
      "                                 (1, 1, 16, 1),                                                   \n",
      "                                 (1, 1, 16, 1),                                                   \n",
      "                                 (1, 1, 16, 1),                                                   \n",
      "                                 (1, 1, 16, 1),                                                   \n",
      "                                 (1, 1, 16, 1),                                                   \n",
      "                                 (1, 1, 16, 1),                                                   \n",
      "                                 (1, 1, 16, 1),                                                   \n",
      "                                 (1, 1, 16, 1),                                                   \n",
      "                                 (1, 1, 16, 1),                                                   \n",
      "                                 (1, 1, 16, 1),                                                   \n",
      "                                 (1, 1, 16, 1),                                                   \n",
      "                                 (1, 1, 16, 1),                                                   \n",
      "                                 (1, 1, 16, 1),                                                   \n",
      "                                 (1, 1, 16, 1),                                                   \n",
      "                                 (1, 1, 16, 1),                                                   \n",
      "                                 (1, 1, 16, 1),                                                   \n",
      "                                 (1, 1, 16, 1),                                                   \n",
      "                                 (1, 1, 16, 1),                                                   \n",
      "                                 (1, 1, 16, 1),                                                   \n",
      "                                 (1, 1, 16, 1),                                                   \n",
      "                                 (1, 1, 16, 1),                                                   \n",
      "                                 (1, 1, 16, 1),                                                   \n",
      "                                 (1, 1, 16, 1),                                                   \n",
      "                                 (1, 1, 16, 1),                                                   \n",
      "                                 (1, 1, 16, 1),                                                   \n",
      "                                 (1, 1, 16, 1),                                                   \n",
      "                                 (1, 1, 16, 1),                                                   \n",
      "                                 (1, 1, 16, 1),                                                   \n",
      "                                 (1, 1, 16, 1),                                                   \n",
      "                                 (1, 1, 16, 1),                                                   \n",
      "                                 (1, 1, 16, 1),                                                   \n",
      "                                 (1, 1, 16, 1),                                                   \n",
      "                                 (1, 1, 16, 1),                                                   \n",
      "                                 (1, 1, 16, 1),                                                   \n",
      "                                 (1, 1, 16, 1),                                                   \n",
      "                                 (1, 1, 16, 1),                                                   \n",
      "                                 (1, 1, 16, 1),                                                   \n",
      "                                 (1, 1, 16, 1),                                                   \n",
      "                                 (1, 1, 16, 1),                                                   \n",
      "                                 (1, 1, 16, 1),                                                   \n",
      "                                 (1, 1, 16, 1),                                                   \n",
      "                                 (1, 1, 16, 1),                                                   \n",
      "                                 (1, 1, 16, 1),                                                   \n",
      "                                 (1, 1, 16, 1),                                                   \n",
      "                                 (1, 1, 16, 1),                                                   \n",
      "                                 (1, 1, 16, 1),                                                   \n",
      "                                 (1, 1, 16, 1),                                                   \n",
      "                                 (1, 1, 16, 1),                                                   \n",
      "                                 (1, 1, 16, 1),                                                   \n",
      "                                 (1, 1, 16, 1),                                                   \n",
      "                                 (1, 1, 16, 1),                                                   \n",
      "                                 (1, 1, 16, 1),                                                   \n",
      "                                 (1, 1, 16, 1),                                                   \n",
      "                                 (1, 1, 16, 1),                                                   \n",
      "                                 (1, 1, 16, 1),                                                   \n",
      "                                 (1, 1, 16, 1),                                                   \n",
      "                                 (1, 1, 16, 1),                                                   \n",
      "                                 (1, 1, 16, 1),                                                   \n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "                                 (1, 1, 16, 1),                                                   \n",
      "                                 (1, 1, 16, 1),                                                   \n",
      "                                 (1, 1, 16, 1)]                                                   \n",
      "                                                                                                  \n",
      " tf.stack_2 (TFOpLambda)        (1, 1, 120, 8)       0           ['GFC_c_0[0][0]',                \n",
      "                                                                  'GFC_c_1[0][0]',                \n",
      "                                                                  'GFC_c_2[0][0]',                \n",
      "                                                                  'GFC_c_3[0][0]',                \n",
      "                                                                  'GFC_c_4[0][0]',                \n",
      "                                                                  'GFC_c_5[0][0]',                \n",
      "                                                                  'GFC_c_6[0][0]',                \n",
      "                                                                  'GFC_c_7[0][0]']                \n",
      "                                                                                                  \n",
      " conv2d_transpose_64 (Conv2DTra  (1, 1, 32, 1)       3           ['tf.split_3[0][0]']             \n",
      " nspose)                                                                                          \n",
      "                                                                                                  \n",
      " conv2d_transpose_65 (Conv2DTra  (1, 1, 32, 1)       3           ['tf.split_3[0][1]']             \n",
      " nspose)                                                                                          \n",
      "                                                                                                  \n",
      " conv2d_transpose_66 (Conv2DTra  (1, 1, 32, 1)       3           ['tf.split_3[0][2]']             \n",
      " nspose)                                                                                          \n",
      "                                                                                                  \n",
      " conv2d_transpose_67 (Conv2DTra  (1, 1, 32, 1)       3           ['tf.split_3[0][3]']             \n",
      " nspose)                                                                                          \n",
      "                                                                                                  \n",
      " conv2d_transpose_68 (Conv2DTra  (1, 1, 32, 1)       3           ['tf.split_3[0][4]']             \n",
      " nspose)                                                                                          \n",
      "                                                                                                  \n",
      " conv2d_transpose_69 (Conv2DTra  (1, 1, 32, 1)       3           ['tf.split_3[0][5]']             \n",
      " nspose)                                                                                          \n",
      "                                                                                                  \n",
      " conv2d_transpose_70 (Conv2DTra  (1, 1, 32, 1)       3           ['tf.split_3[0][6]']             \n",
      " nspose)                                                                                          \n",
      "                                                                                                  \n",
      " conv2d_transpose_71 (Conv2DTra  (1, 1, 32, 1)       3           ['tf.split_3[0][7]']             \n",
      " nspose)                                                                                          \n",
      "                                                                                                  \n",
      " conv2d_transpose_72 (Conv2DTra  (1, 1, 32, 1)       3           ['tf.split_3[0][8]']             \n",
      " nspose)                                                                                          \n",
      "                                                                                                  \n",
      " conv2d_transpose_73 (Conv2DTra  (1, 1, 32, 1)       3           ['tf.split_3[0][9]']             \n",
      " nspose)                                                                                          \n",
      "                                                                                                  \n",
      " conv2d_transpose_74 (Conv2DTra  (1, 1, 32, 1)       3           ['tf.split_3[0][10]']            \n",
      " nspose)                                                                                          \n",
      "                                                                                                  \n",
      " conv2d_transpose_75 (Conv2DTra  (1, 1, 32, 1)       3           ['tf.split_3[0][11]']            \n",
      " nspose)                                                                                          \n",
      "                                                                                                  \n",
      " conv2d_transpose_76 (Conv2DTra  (1, 1, 32, 1)       3           ['tf.split_3[0][12]']            \n",
      " nspose)                                                                                          \n",
      "                                                                                                  \n",
      " conv2d_transpose_77 (Conv2DTra  (1, 1, 32, 1)       3           ['tf.split_3[0][13]']            \n",
      " nspose)                                                                                          \n",
      "                                                                                                  \n",
      " conv2d_transpose_78 (Conv2DTra  (1, 1, 32, 1)       3           ['tf.split_3[0][14]']            \n",
      " nspose)                                                                                          \n",
      "                                                                                                  \n",
      " conv2d_transpose_79 (Conv2DTra  (1, 1, 32, 1)       3           ['tf.split_3[0][15]']            \n",
      " nspose)                                                                                          \n",
      "                                                                                                  \n",
      " conv2d_transpose_80 (Conv2DTra  (1, 1, 32, 1)       3           ['tf.split_3[0][16]']            \n",
      " nspose)                                                                                          \n",
      "                                                                                                  \n",
      " conv2d_transpose_81 (Conv2DTra  (1, 1, 32, 1)       3           ['tf.split_3[0][17]']            \n",
      " nspose)                                                                                          \n",
      "                                                                                                  \n",
      " conv2d_transpose_82 (Conv2DTra  (1, 1, 32, 1)       3           ['tf.split_3[0][18]']            \n",
      " nspose)                                                                                          \n",
      "                                                                                                  \n",
      " conv2d_transpose_83 (Conv2DTra  (1, 1, 32, 1)       3           ['tf.split_3[0][19]']            \n",
      " nspose)                                                                                          \n",
      "                                                                                                  \n",
      " conv2d_transpose_84 (Conv2DTra  (1, 1, 32, 1)       3           ['tf.split_3[0][20]']            \n",
      " nspose)                                                                                          \n",
      "                                                                                                  \n",
      " conv2d_transpose_85 (Conv2DTra  (1, 1, 32, 1)       3           ['tf.split_3[0][21]']            \n",
      " nspose)                                                                                          \n",
      "                                                                                                  \n",
      " conv2d_transpose_86 (Conv2DTra  (1, 1, 32, 1)       3           ['tf.split_3[0][22]']            \n",
      " nspose)                                                                                          \n",
      "                                                                                                  \n",
      " conv2d_transpose_87 (Conv2DTra  (1, 1, 32, 1)       3           ['tf.split_3[0][23]']            \n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      " nspose)                                                                                          \n",
      "                                                                                                  \n",
      " conv2d_transpose_88 (Conv2DTra  (1, 1, 32, 1)       3           ['tf.split_3[0][24]']            \n",
      " nspose)                                                                                          \n",
      "                                                                                                  \n",
      " conv2d_transpose_89 (Conv2DTra  (1, 1, 32, 1)       3           ['tf.split_3[0][25]']            \n",
      " nspose)                                                                                          \n",
      "                                                                                                  \n",
      " conv2d_transpose_90 (Conv2DTra  (1, 1, 32, 1)       3           ['tf.split_3[0][26]']            \n",
      " nspose)                                                                                          \n",
      "                                                                                                  \n",
      " conv2d_transpose_91 (Conv2DTra  (1, 1, 32, 1)       3           ['tf.split_3[0][27]']            \n",
      " nspose)                                                                                          \n",
      "                                                                                                  \n",
      " conv2d_transpose_92 (Conv2DTra  (1, 1, 32, 1)       3           ['tf.split_3[0][28]']            \n",
      " nspose)                                                                                          \n",
      "                                                                                                  \n",
      " conv2d_transpose_93 (Conv2DTra  (1, 1, 32, 1)       3           ['tf.split_3[0][29]']            \n",
      " nspose)                                                                                          \n",
      "                                                                                                  \n",
      " conv2d_transpose_94 (Conv2DTra  (1, 1, 32, 1)       3           ['tf.split_3[0][30]']            \n",
      " nspose)                                                                                          \n",
      "                                                                                                  \n",
      " conv2d_transpose_95 (Conv2DTra  (1, 1, 32, 1)       3           ['tf.split_3[0][31]']            \n",
      " nspose)                                                                                          \n",
      "                                                                                                  \n",
      " conv2d_transpose_96 (Conv2DTra  (1, 1, 32, 1)       3           ['tf.split_3[0][32]']            \n",
      " nspose)                                                                                          \n",
      "                                                                                                  \n",
      " conv2d_transpose_97 (Conv2DTra  (1, 1, 32, 1)       3           ['tf.split_3[0][33]']            \n",
      " nspose)                                                                                          \n",
      "                                                                                                  \n",
      " conv2d_transpose_98 (Conv2DTra  (1, 1, 32, 1)       3           ['tf.split_3[0][34]']            \n",
      " nspose)                                                                                          \n",
      "                                                                                                  \n",
      " conv2d_transpose_99 (Conv2DTra  (1, 1, 32, 1)       3           ['tf.split_3[0][35]']            \n",
      " nspose)                                                                                          \n",
      "                                                                                                  \n",
      " conv2d_transpose_100 (Conv2DTr  (1, 1, 32, 1)       3           ['tf.split_3[0][36]']            \n",
      " anspose)                                                                                         \n",
      "                                                                                                  \n",
      " conv2d_transpose_101 (Conv2DTr  (1, 1, 32, 1)       3           ['tf.split_3[0][37]']            \n",
      " anspose)                                                                                         \n",
      "                                                                                                  \n",
      " conv2d_transpose_102 (Conv2DTr  (1, 1, 32, 1)       3           ['tf.split_3[0][38]']            \n",
      " anspose)                                                                                         \n",
      "                                                                                                  \n",
      " conv2d_transpose_103 (Conv2DTr  (1, 1, 32, 1)       3           ['tf.split_3[0][39]']            \n",
      " anspose)                                                                                         \n",
      "                                                                                                  \n",
      " conv2d_transpose_104 (Conv2DTr  (1, 1, 32, 1)       3           ['tf.split_3[0][40]']            \n",
      " anspose)                                                                                         \n",
      "                                                                                                  \n",
      " conv2d_transpose_105 (Conv2DTr  (1, 1, 32, 1)       3           ['tf.split_3[0][41]']            \n",
      " anspose)                                                                                         \n",
      "                                                                                                  \n",
      " conv2d_transpose_106 (Conv2DTr  (1, 1, 32, 1)       3           ['tf.split_3[0][42]']            \n",
      " anspose)                                                                                         \n",
      "                                                                                                  \n",
      " conv2d_transpose_107 (Conv2DTr  (1, 1, 32, 1)       3           ['tf.split_3[0][43]']            \n",
      " anspose)                                                                                         \n",
      "                                                                                                  \n",
      " conv2d_transpose_108 (Conv2DTr  (1, 1, 32, 1)       3           ['tf.split_3[0][44]']            \n",
      " anspose)                                                                                         \n",
      "                                                                                                  \n",
      " conv2d_transpose_109 (Conv2DTr  (1, 1, 32, 1)       3           ['tf.split_3[0][45]']            \n",
      " anspose)                                                                                         \n",
      "                                                                                                  \n",
      " conv2d_transpose_110 (Conv2DTr  (1, 1, 32, 1)       3           ['tf.split_3[0][46]']            \n",
      " anspose)                                                                                         \n",
      "                                                                                                  \n",
      " conv2d_transpose_111 (Conv2DTr  (1, 1, 32, 1)       3           ['tf.split_3[0][47]']            \n",
      " anspose)                                                                                         \n",
      "                                                                                                  \n",
      " conv2d_transpose_112 (Conv2DTr  (1, 1, 32, 1)       3           ['tf.split_3[0][48]']            \n",
      " anspose)                                                                                         \n",
      "                                                                                                  \n",
      " conv2d_transpose_113 (Conv2DTr  (1, 1, 32, 1)       3           ['tf.split_3[0][49]']            \n",
      " anspose)                                                                                         \n",
      "                                                                                                  \n",
      " conv2d_transpose_114 (Conv2DTr  (1, 1, 32, 1)       3           ['tf.split_3[0][50]']            \n",
      " anspose)                                                                                         \n",
      "                                                                                                  \n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      " conv2d_transpose_115 (Conv2DTr  (1, 1, 32, 1)       3           ['tf.split_3[0][51]']            \n",
      " anspose)                                                                                         \n",
      "                                                                                                  \n",
      " conv2d_transpose_116 (Conv2DTr  (1, 1, 32, 1)       3           ['tf.split_3[0][52]']            \n",
      " anspose)                                                                                         \n",
      "                                                                                                  \n",
      " conv2d_transpose_117 (Conv2DTr  (1, 1, 32, 1)       3           ['tf.split_3[0][53]']            \n",
      " anspose)                                                                                         \n",
      "                                                                                                  \n",
      " conv2d_transpose_118 (Conv2DTr  (1, 1, 32, 1)       3           ['tf.split_3[0][54]']            \n",
      " anspose)                                                                                         \n",
      "                                                                                                  \n",
      " conv2d_transpose_119 (Conv2DTr  (1, 1, 32, 1)       3           ['tf.split_3[0][55]']            \n",
      " anspose)                                                                                         \n",
      "                                                                                                  \n",
      " conv2d_transpose_120 (Conv2DTr  (1, 1, 32, 1)       3           ['tf.split_3[0][56]']            \n",
      " anspose)                                                                                         \n",
      "                                                                                                  \n",
      " conv2d_transpose_121 (Conv2DTr  (1, 1, 32, 1)       3           ['tf.split_3[0][57]']            \n",
      " anspose)                                                                                         \n",
      "                                                                                                  \n",
      " conv2d_transpose_122 (Conv2DTr  (1, 1, 32, 1)       3           ['tf.split_3[0][58]']            \n",
      " anspose)                                                                                         \n",
      "                                                                                                  \n",
      " conv2d_transpose_123 (Conv2DTr  (1, 1, 32, 1)       3           ['tf.split_3[0][59]']            \n",
      " anspose)                                                                                         \n",
      "                                                                                                  \n",
      " conv2d_transpose_124 (Conv2DTr  (1, 1, 32, 1)       3           ['tf.split_3[0][60]']            \n",
      " anspose)                                                                                         \n",
      "                                                                                                  \n",
      " conv2d_transpose_125 (Conv2DTr  (1, 1, 32, 1)       3           ['tf.split_3[0][61]']            \n",
      " anspose)                                                                                         \n",
      "                                                                                                  \n",
      " conv2d_transpose_126 (Conv2DTr  (1, 1, 32, 1)       3           ['tf.split_3[0][62]']            \n",
      " anspose)                                                                                         \n",
      "                                                                                                  \n",
      " conv2d_transpose_127 (Conv2DTr  (1, 1, 32, 1)       3           ['tf.split_3[0][63]']            \n",
      " anspose)                                                                                         \n",
      "                                                                                                  \n",
      " tf.compat.v1.transpose_5 (TFOp  (1, 1, 8, 120)      0           ['tf.stack_2[0][0]']             \n",
      " Lambda)                                                                                          \n",
      "                                                                                                  \n",
      " tf.concat_1 (TFOpLambda)       (1, 1, 32, 64)       0           ['conv2d_transpose_64[0][0]',    \n",
      "                                                                  'conv2d_transpose_65[0][0]',    \n",
      "                                                                  'conv2d_transpose_66[0][0]',    \n",
      "                                                                  'conv2d_transpose_67[0][0]',    \n",
      "                                                                  'conv2d_transpose_68[0][0]',    \n",
      "                                                                  'conv2d_transpose_69[0][0]',    \n",
      "                                                                  'conv2d_transpose_70[0][0]',    \n",
      "                                                                  'conv2d_transpose_71[0][0]',    \n",
      "                                                                  'conv2d_transpose_72[0][0]',    \n",
      "                                                                  'conv2d_transpose_73[0][0]',    \n",
      "                                                                  'conv2d_transpose_74[0][0]',    \n",
      "                                                                  'conv2d_transpose_75[0][0]',    \n",
      "                                                                  'conv2d_transpose_76[0][0]',    \n",
      "                                                                  'conv2d_transpose_77[0][0]',    \n",
      "                                                                  'conv2d_transpose_78[0][0]',    \n",
      "                                                                  'conv2d_transpose_79[0][0]',    \n",
      "                                                                  'conv2d_transpose_80[0][0]',    \n",
      "                                                                  'conv2d_transpose_81[0][0]',    \n",
      "                                                                  'conv2d_transpose_82[0][0]',    \n",
      "                                                                  'conv2d_transpose_83[0][0]',    \n",
      "                                                                  'conv2d_transpose_84[0][0]',    \n",
      "                                                                  'conv2d_transpose_85[0][0]',    \n",
      "                                                                  'conv2d_transpose_86[0][0]',    \n",
      "                                                                  'conv2d_transpose_87[0][0]',    \n",
      "                                                                  'conv2d_transpose_88[0][0]',    \n",
      "                                                                  'conv2d_transpose_89[0][0]',    \n",
      "                                                                  'conv2d_transpose_90[0][0]',    \n",
      "                                                                  'conv2d_transpose_91[0][0]',    \n",
      "                                                                  'conv2d_transpose_92[0][0]',    \n",
      "                                                                  'conv2d_transpose_93[0][0]',    \n",
      "                                                                  'conv2d_transpose_94[0][0]',    \n",
      "                                                                  'conv2d_transpose_95[0][0]',    \n",
      "                                                                  'conv2d_transpose_96[0][0]',    \n",
      "                                                                  'conv2d_transpose_97[0][0]',    \n",
      "                                                                  'conv2d_transpose_98[0][0]',    \n",
      "                                                                  'conv2d_transpose_99[0][0]',    \n",
      "                                                                  'conv2d_transpose_100[0][0]',   \n",
      "                                                                  'conv2d_transpose_101[0][0]',   \n",
      "                                                                  'conv2d_transpose_102[0][0]',   \n",
      "                                                                  'conv2d_transpose_103[0][0]',   \n",
      "                                                                  'conv2d_transpose_104[0][0]',   \n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "                                                                  'conv2d_transpose_105[0][0]',   \n",
      "                                                                  'conv2d_transpose_106[0][0]',   \n",
      "                                                                  'conv2d_transpose_107[0][0]',   \n",
      "                                                                  'conv2d_transpose_108[0][0]',   \n",
      "                                                                  'conv2d_transpose_109[0][0]',   \n",
      "                                                                  'conv2d_transpose_110[0][0]',   \n",
      "                                                                  'conv2d_transpose_111[0][0]',   \n",
      "                                                                  'conv2d_transpose_112[0][0]',   \n",
      "                                                                  'conv2d_transpose_113[0][0]',   \n",
      "                                                                  'conv2d_transpose_114[0][0]',   \n",
      "                                                                  'conv2d_transpose_115[0][0]',   \n",
      "                                                                  'conv2d_transpose_116[0][0]',   \n",
      "                                                                  'conv2d_transpose_117[0][0]',   \n",
      "                                                                  'conv2d_transpose_118[0][0]',   \n",
      "                                                                  'conv2d_transpose_119[0][0]',   \n",
      "                                                                  'conv2d_transpose_120[0][0]',   \n",
      "                                                                  'conv2d_transpose_121[0][0]',   \n",
      "                                                                  'conv2d_transpose_122[0][0]',   \n",
      "                                                                  'conv2d_transpose_123[0][0]',   \n",
      "                                                                  'conv2d_transpose_124[0][0]',   \n",
      "                                                                  'conv2d_transpose_125[0][0]',   \n",
      "                                                                  'conv2d_transpose_126[0][0]',   \n",
      "                                                                  'conv2d_transpose_127[0][0]']   \n",
      "                                                                                                  \n",
      " tf.compat.v1.shape_27 (TFOpLam  (4,)                0           ['tf.compat.v1.transpose_5[0][0]'\n",
      " bda)                                                            ]                                \n",
      "                                                                                                  \n",
      " tf.compat.v1.shape_26 (TFOpLam  (4,)                0           ['tf.compat.v1.transpose_5[0][0]'\n",
      " bda)                                                            ]                                \n",
      "                                                                                                  \n",
      " convp0 (Conv2D)                (1, 1, 32, 64)       64          ['tf.nn.relu[0][0]']             \n",
      "                                                                                                  \n",
      " convt1_1x1 (Conv2D)            (1, 1, 32, 64)       4096        ['tf.concat_1[0][0]']            \n",
      "                                                                                                  \n",
      " tf.compat.v1.shape_24 (TFOpLam  (4,)                0           ['tf.compat.v1.transpose_5[0][0]'\n",
      " bda)                                                            ]                                \n",
      "                                                                                                  \n",
      " tf.compat.v1.shape_25 (TFOpLam  (4,)                0           ['tf.compat.v1.transpose_5[0][0]'\n",
      " bda)                                                            ]                                \n",
      "                                                                                                  \n",
      " tf.__operators__.getitem_27 (S  ()                  0           ['tf.compat.v1.shape_27[0][0]']  \n",
      " licingOpLambda)                                                                                  \n",
      "                                                                                                  \n",
      " tf.__operators__.getitem_26 (S  ()                  0           ['tf.compat.v1.shape_26[0][0]']  \n",
      " licingOpLambda)                                                                                  \n",
      "                                                                                                  \n",
      " convp0BatchNorm (BatchNormaliz  (1, 1, 32, 64)      256         ['convp0[0][0]']                 \n",
      " ation)                                                                                           \n",
      "                                                                                                  \n",
      " convt1BatchNorm (BatchNormaliz  (1, 1, 32, 64)      256         ['convt1_1x1[0][0]']             \n",
      " ation)                                                                                           \n",
      "                                                                                                  \n",
      " tf.__operators__.getitem_24 (S  ()                  0           ['tf.compat.v1.shape_24[0][0]']  \n",
      " licingOpLambda)                                                                                  \n",
      "                                                                                                  \n",
      " tf.__operators__.getitem_25 (S  ()                  0           ['tf.compat.v1.shape_25[0][0]']  \n",
      " licingOpLambda)                                                                                  \n",
      "                                                                                                  \n",
      " tf.math.multiply_4 (TFOpLambda  ()                  0           ['tf.__operators__.getitem_27[0][\n",
      " )                                                               0]',                             \n",
      "                                                                  'tf.__operators__.getitem_26[0][\n",
      "                                                                 0]']                             \n",
      "                                                                                                  \n",
      " tf.nn.relu_16 (TFOpLambda)     (1, 1, 32, 64)       0           ['convp0BatchNorm[0][0]']        \n",
      "                                                                                                  \n",
      " tf.nn.relu_15 (TFOpLambda)     (1, 1, 32, 64)       0           ['convt1BatchNorm[0][0]']        \n",
      "                                                                                                  \n",
      " tf.reshape_6 (TFOpLambda)      (1, 1, 960)          0           ['tf.compat.v1.transpose_5[0][0]'\n",
      "                                                                 , 'tf.__operators__.getitem_24[0]\n",
      "                                                                 [0]',                            \n",
      "                                                                  'tf.__operators__.getitem_25[0][\n",
      "                                                                 0]',                             \n",
      "                                                                  'tf.math.multiply_4[0][0]']     \n",
      "                                                                                                  \n",
      " tf.__operators__.add_4 (TFOpLa  (1, 1, 32, 64)      0           ['tf.nn.relu_16[0][0]',          \n",
      " mbda)                                                            'tf.nn.relu_15[0][0]']          \n",
      "                                                                                                  \n",
      " tf.math.tanh (TFOpLambda)      (1, 1, 960)          0           ['tf.reshape_6[0][0]']           \n",
      "                                                                                                  \n",
      " convp_DfDecoder (Conv2D)       (1, 1, 96, 10)       320         ['tf.nn.relu_4[0][0]']           \n",
      "                                                                                                  \n",
      " tf.compat.v1.pad_7 (TFOpLambda  (1, 1, 34, 64)      0           ['tf.__operators__.add_4[0][0]'] \n",
      " )                                                                                                \n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "                                                                                                  \n",
      " tf.compat.v1.shape_28 (TFOpLam  (3,)                0           ['tf.math.tanh[0][0]']           \n",
      " bda)                                                                                             \n",
      "                                                                                                  \n",
      " tf.compat.v1.shape_29 (TFOpLam  (3,)                0           ['tf.math.tanh[0][0]']           \n",
      " bda)                                                                                             \n",
      "                                                                                                  \n",
      " convp_DfDecoderBatchNorm (Batc  (1, 1, 96, 10)      40          ['convp_DfDecoder[0][0]']        \n",
      " hNormalization)                                                                                  \n",
      "                                                                                                  \n",
      " mask_out (Conv2D)              (1, 1, 32, 1)        192         ['tf.compat.v1.pad_7[0][0]']     \n",
      "                                                                                                  \n",
      " tf.__operators__.getitem_28 (S  ()                  0           ['tf.compat.v1.shape_28[0][0]']  \n",
      " licingOpLambda)                                                                                  \n",
      "                                                                                                  \n",
      " tf.__operators__.getitem_29 (S  ()                  0           ['tf.compat.v1.shape_29[0][0]']  \n",
      " licingOpLambda)                                                                                  \n",
      "                                                                                                  \n",
      " tf.nn.relu_19 (TFOpLambda)     (1, 1, 96, 10)       0           ['convp_DfDecoderBatchNorm[0][0]'\n",
      "                                                                 ]                                \n",
      "                                                                                                  \n",
      " tf.compat.v1.shape_20 (TFOpLam  (4,)                0           ['mask_out[0][0]']               \n",
      " bda)                                                                                             \n",
      "                                                                                                  \n",
      " tf.compat.v1.shape_21 (TFOpLam  (4,)                0           ['mask_out[0][0]']               \n",
      " bda)                                                                                             \n",
      "                                                                                                  \n",
      " tf.compat.v1.shape_22 (TFOpLam  (4,)                0           ['mask_out[0][0]']               \n",
      " bda)                                                                                             \n",
      "                                                                                                  \n",
      " tf.reshape_7 (TFOpLambda)      (1, 1, 10, 96)       0           ['tf.math.tanh[0][0]',           \n",
      "                                                                  'tf.__operators__.getitem_28[0][\n",
      "                                                                 0]',                             \n",
      "                                                                  'tf.__operators__.getitem_29[0][\n",
      "                                                                 0]']                             \n",
      "                                                                                                  \n",
      " tf.compat.v1.transpose_4 (TFOp  (1, 1, 10, 96)      0           ['tf.nn.relu_19[0][0]']          \n",
      " Lambda)                                                                                          \n",
      "                                                                                                  \n",
      " tf.__operators__.getitem_20 (S  ()                  0           ['tf.compat.v1.shape_20[0][0]']  \n",
      " licingOpLambda)                                                                                  \n",
      "                                                                                                  \n",
      " tf.__operators__.getitem_21 (S  ()                  0           ['tf.compat.v1.shape_21[0][0]']  \n",
      " licingOpLambda)                                                                                  \n",
      "                                                                                                  \n",
      " tf.__operators__.getitem_22 (S  ()                  0           ['tf.compat.v1.shape_22[0][0]']  \n",
      " licingOpLambda)                                                                                  \n",
      "                                                                                                  \n",
      " tf.__operators__.add_6 (TFOpLa  (1, 1, 10, 96)      0           ['tf.reshape_7[0][0]',           \n",
      " mbda)                                                            'tf.compat.v1.transpose_4[0][0]'\n",
      "                                                                 ]                                \n",
      "                                                                                                  \n",
      " tf.reshape_5 (TFOpLambda)      (1, 1, 32)           0           ['mask_out[0][0]',               \n",
      "                                                                  'tf.__operators__.getitem_20[0][\n",
      "                                                                 0]',                             \n",
      "                                                                  'tf.__operators__.getitem_21[0][\n",
      "                                                                 0]',                             \n",
      "                                                                  'tf.__operators__.getitem_22[0][\n",
      "                                                                 0]']                             \n",
      "                                                                                                  \n",
      " tf.reshape_8 (TFOpLambda)      (1, 1, 5, 2, 96)     0           ['tf.__operators__.add_6[0][0]', \n",
      "                                                                  'tf.__operators__.getitem_28[0][\n",
      "                                                                 0]',                             \n",
      "                                                                  'tf.__operators__.getitem_29[0][\n",
      "                                                                 0]']                             \n",
      "                                                                                                  \n",
      " tf.math.sigmoid (TFOpLambda)   (1, 1, 32)           0           ['tf.reshape_5[0][0]']           \n",
      "                                                                                                  \n",
      " convp_alpha (Dense)            (1, 1, 1)            257         ['tf.__operators__.add_5[0][0]'] \n",
      "                                                                                                  \n",
      " tf.compat.v1.transpose_6 (TFOp  (1, 1, 5, 96, 2)    0           ['tf.reshape_8[0][0]']           \n",
      " Lambda)                                                                                          \n",
      "                                                                                                  \n",
      "==================================================================================================\n",
      "Total params: 2,848,937\n",
      "Trainable params: 0\n",
      "Non-trainable params: 2,848,937\n",
      "__________________________________________________________________________________________________\n"
     ]
    }
   ],
   "source": [
    "model.summary()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "id": "d9d31432",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "WARNING:tensorflow:From /home/myhsueh/.local/lib/python3.8/site-packages/tensorflow/python/ops/nn_ops.py:5219: tensor_shape_from_node_def_name (from tensorflow.python.framework.graph_util_impl) is deprecated and will be removed in a future version.\n",
      "Instructions for updating:\n",
      "Use `tf.compat.v1.graph_util.tensor_shape_from_node_def_name`\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "WARNING:tensorflow:From /home/myhsueh/.local/lib/python3.8/site-packages/tensorflow/python/ops/nn_ops.py:5219: tensor_shape_from_node_def_name (from tensorflow.python.framework.graph_util_impl) is deprecated and will be removed in a future version.\n",
      "Instructions for updating:\n",
      "Use `tf.compat.v1.graph_util.tensor_shape_from_node_def_name`\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "5.651398 M Flops\n",
      "2.825699 M MACS\n"
     ]
    }
   ],
   "source": [
    "from typing import Any, Callable, Dict, List, Optional, Union\n",
    "from tensorflow.python.framework.convert_to_constants import convert_variables_to_constants_v2_as_graph\n",
    "\n",
    "def try_count_flops(model: Union[tf.Module, tf.keras.Model],\n",
    "                    inputs_kwargs: Optional[Dict[str, Any]] = None,\n",
    "                    output_path: Optional[str] = None):\n",
    "    \"\"\"Counts and returns model FLOPs.\n",
    "  Args:\n",
    "    model: A model instance.\n",
    "    inputs_kwargs: An optional dictionary of argument pairs specifying inputs'\n",
    "      shape specifications to getting corresponding concrete function.\n",
    "    output_path: A file path to write the profiling results to.\n",
    "  Returns:\n",
    "    The model's FLOPs.\n",
    "  \"\"\"\n",
    "    if hasattr(model, 'inputs'):\n",
    "        try:\n",
    "            # Get input shape and set batch size to 1.\n",
    "            if model.inputs:\n",
    "                inputs = [\n",
    "                    tf.TensorSpec([1] + input.shape[1:], input.dtype)\n",
    "                    for input in model.inputs\n",
    "                ]\n",
    "                concrete_func = tf.function(model).get_concrete_function(inputs)\n",
    "            # If model.inputs is invalid, try to use the input to get concrete\n",
    "            # function for model.call (subclass model).\n",
    "            else:\n",
    "                concrete_func = tf.function(model.call).get_concrete_function(\n",
    "                    **inputs_kwargs)\n",
    "            frozen_func, _ = convert_variables_to_constants_v2_as_graph(concrete_func)\n",
    "\n",
    "            # Calculate FLOPs.\n",
    "            run_meta = tf.compat.v1.RunMetadata()\n",
    "            opts = tf.compat.v1.profiler.ProfileOptionBuilder.float_operation()\n",
    "            if output_path is not None:\n",
    "                opts['output'] = f'file:outfile={output_path}'\n",
    "            else:\n",
    "                opts['output'] = 'none'\n",
    "            flops = tf.compat.v1.profiler.profile(\n",
    "                graph=frozen_func.graph, run_meta=run_meta, options=opts)\n",
    "            return flops.total_float_ops\n",
    "        except Exception as e:  # pylint: disable=broad-except\n",
    "            logging.info(\n",
    "                'Failed to count model FLOPs with error %s, because the build() '\n",
    "                'methods in keras layers were not called. This is probably because '\n",
    "                'the model was not feed any input, e.g., the max train step already '\n",
    "                'reached before this run.', e)\n",
    "            return None\n",
    "    return None\n",
    "\n",
    "flops = try_count_flops(model)\n",
    "print(flops/1000000,\"M Flops\")\n",
    "print(flops/2000000,\"M MACS\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "500705ec",
   "metadata": {},
   "source": [
    "# Pre-train weight"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "id": "da1ecac0",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "weights.19-58.49.h5\n"
     ]
    }
   ],
   "source": [
    "# path = \"./weights8/\"\n",
    "# h5_lists = os.listdir(path)\n",
    "# h5_lists.sort(key=lambda fn:os.path.getmtime(path + fn)\n",
    "#                 if not os.path.isdir(path + fn) else 0)\n",
    "# new_file = os.path.join(path, h5_lists[-1])\n",
    "new_file = 'weights.19-58.49.h5'\n",
    "print(new_file)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "id": "7abf1686",
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/home/myhsueh/.local/lib/python3.8/site-packages/numpy/core/numeric.py:2443: VisibleDeprecationWarning: Creating an ndarray from ragged nested sequences (which is a list-or-tuple of lists-or-tuples-or ndarrays with different lengths or shapes) is deprecated. If you meant to do this, you must specify 'dtype=object' when creating the ndarray.\n",
      "  a1, a2 = asarray(a1), asarray(a2)\n"
     ]
    }
   ],
   "source": [
    "# store weights before loading pre-trained weights\n",
    "preloaded_layers = model.layers.copy()\n",
    "preloaded_weights = []\n",
    "for pre in preloaded_layers:\n",
    "    preloaded_weights.append(pre.get_weights())\n",
    "\n",
    "# load pre-trained weights\n",
    "model.load_weights(new_file, by_name=True)\n",
    "\n",
    "# compare previews weights vs loaded weights\n",
    "for layer, pre in zip(model.layers, preloaded_weights):\n",
    "    weights = layer.get_weights()\n",
    "\n",
    "    if weights:\n",
    "        if np.array_equal(weights, pre):\n",
    "            print('not loaded', layer.name)\n",
    "        else:\n",
    "#             print('loaded', layer.name)\n",
    "            pass"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7c12ec5d",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "id": "b589b723",
   "metadata": {},
   "outputs": [],
   "source": [
    "def band_unit_norm(x, s, alpha=0.99):\n",
    "    s = np.linalg.norm(x) * (1-alpha) + s * alpha\n",
    "    x = x / (np.sqrt(s)+1e-12)\n",
    "    return x, s\n",
    "def band_mean_norm_erb(x, s, alpha=0.99):\n",
    "    s = x * (1-alpha) + s * alpha\n",
    "    x = (x-s) / 40\n",
    "    return x, s\n",
    "\n",
    "def erb_norm(x, mean_init=[-60.0,-90.0]):\n",
    "    # x : [T,F,C]\n",
    "    # state: [F,C]\n",
    "    x = np.reshape(x,(x.shape[0],x.shape[1],1))\n",
    "    shape = x.shape\n",
    "    state = np.linspace(mean_init[0],mean_init[1],shape[-2]) # [F,]\n",
    "#     state *= 0\n",
    "    state = np.reshape(state, (1, shape[-2]))\n",
    "    state = np.tile(state, (1, shape[-1])) # [C,F]\n",
    "    \n",
    "    x_i_list = []\n",
    "    state_list = []\n",
    "    \n",
    "    for i in range(shape[-1]):\n",
    "        x_i = np.split(x, shape[-1], axis=-1) \n",
    "        state_i = np.split(state, shape[-1], axis=-1)\n",
    "        \n",
    "        x_ij_list = []\n",
    "        state_tmp = state_i[i]\n",
    "        for j in range(shape[-3]):\n",
    "            x_ij = np.split(x_i[i], shape[0], axis=0) \n",
    "            x_tmp, state_tmp = band_mean_norm_erb(np.squeeze(x_ij[j],-1), state_tmp)\n",
    "            x_ij_list.append(x_tmp)\n",
    "        \n",
    "        x_i_list.append(np.stack(x_ij_list,1))\n",
    "        state_list.append(state_tmp)\n",
    "    x = np.squeeze(np.stack(x_i_list,-1),0)\n",
    "    return x\n",
    "\n",
    "def unit_norm(x, unit_init=[0.001, 0.0001]):\n",
    "    # x : [T,F,C]\n",
    "    # state: [F,C]\n",
    "    x = np.reshape(x,(x.shape[0],x.shape[1],1))\n",
    "    shape = x.shape\n",
    "    state = np.linspace(unit_init[0],unit_init[1],shape[-2]) # [F,]\n",
    "#     state *= 0\n",
    "    state = np.reshape(state, (1, shape[-2]))\n",
    "    state = np.tile(state, (1, shape[-1])) # [C,F]\n",
    "    \n",
    "    x_i_list = []\n",
    "    state_list = []\n",
    "    \n",
    "    for i in range(shape[-1]):\n",
    "        x_i = np.split(x, shape[-1], axis=-1) \n",
    "        state_i = np.split(state, shape[-1], axis=-1)\n",
    "        \n",
    "        x_ij_list = []\n",
    "        state_tmp = state_i[i]\n",
    "        for j in range(shape[-3]):\n",
    "            x_ij = np.split(x_i[i], shape[0], axis=0) \n",
    "            x_tmp, state_tmp = band_unit_norm(np.squeeze(x_ij[j],-1), state_tmp)\n",
    "            x_ij_list.append(x_tmp)\n",
    "        \n",
    "        x_i_list.append(np.stack(x_ij_list,1))\n",
    "        state_list.append(state_tmp)\n",
    "    \n",
    "    x = np.squeeze(np.stack(x_i_list,-1))\n",
    "#     print('done spec')\n",
    "    return x"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "id": "393e9328",
   "metadata": {},
   "outputs": [],
   "source": [
    "ERBB = ERBBand(N=p.nb_erb, high_lim=p.sr//2, NFFT=p.fft_size)\n",
    "ERB_Matrix = ERB_pro_matrix(ERBB, NFFT=p.fft_size, mode=0)\n",
    "iERB_Matrix = ERB_pro_matrix(ERBB, NFFT=p.fft_size, mode=1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "id": "f9c998a0",
   "metadata": {},
   "outputs": [],
   "source": [
    "import librosa\n",
    "from soundfile import write\n",
    "import matplotlib.pyplot as plt"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "id": "159b5377",
   "metadata": {},
   "outputs": [],
   "source": [
    "# import tensorflow as tf\n",
    "def vorbis_window(FRAME_SIZE, transpose=True):\n",
    "    FRAME_SIZE = FRAME_SIZE//2\n",
    "    win = np.zeros((FRAME_SIZE,))\n",
    "    for i in range(FRAME_SIZE):\n",
    "        win[i] = np.sin(.5*np.pi*np.sin(.5*np.pi*(i+.5)/FRAME_SIZE) * np.sin(.5*np.pi*(i+.5)/FRAME_SIZE))\n",
    "\n",
    "    win = np.concatenate((win,np.flip(win)),0)\n",
    "    if transpose: win = win.T\n",
    "    return win\n",
    "\n",
    "fft_norm = (p.fft_size ** -0.5)\n",
    "\n",
    "def analysis_frame(x, nfft=p.fft_size, hop=p.hop_size, normalize=False):\n",
    "    length = len(x)\n",
    "    n_frames = length // hop\n",
    "    out = np.empty((n_frames, p.fft_size//2+1),dtype=complex)\n",
    "    if not length % hop == 0:\n",
    "        x = np.pad(x,(0, nfft - length%hop))\n",
    "    for frame_idx in range(0, n_frames * hop, hop):\n",
    "        frame = x[frame_idx : frame_idx + nfft]\n",
    "        if len(frame)<nfft: frame = np.pad(frame,(0,nfft-len(frame)))\n",
    "        win = vorbis_window(nfft)\n",
    "        frame = frame.reshape(win.shape)\n",
    "        frame_win = np.multiply(frame, win)\n",
    "        x_fft = np.fft.rfft(frame_win, n=p.fft_size) \n",
    "        if normalize: x_fft *= fft_norm\n",
    "        out[frame_idx//hop,:] = x_fft\n",
    "    return out \n",
    "\n",
    "def synthesis_frame(x, nfft=p.fft_size, hop=p.hop_size, windowing=False, normalize=False):\n",
    "    n_frames, _ = x.shape\n",
    "    length = (n_frames+1)*hop\n",
    "    out = np.zeros((length,))\n",
    "#     x[:,:2] = x[:,:2]*0.01\n",
    "    win = vorbis_window(nfft, transpose=False)\n",
    "    if normalize: x /= fft_norm\n",
    "    for frame_idx in range(0, n_frames * hop, hop):\n",
    "        frame = np.fft.irfft(x[frame_idx//hop] , n=p.fft_size)\n",
    "        if windowing:\n",
    "            frame = frame.reshape(win.shape)\n",
    "            out[frame_idx : frame_idx + nfft] += np.multiply(frame, win)\n",
    "        else:\n",
    "            out[frame_idx : frame_idx + nfft] += frame\n",
    "    return out"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "df272655",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "id": "a00698a2",
   "metadata": {},
   "outputs": [],
   "source": [
    "def process(audiofile, process_path='/home/myhsueh/DeepFilterNet/wav/proc/dfn2_sgd/gain/'):\n",
    "    y, sr = librosa.load(audiofile, sr=None)\n",
    "    if not sr==p.sr: y = librosa.resample(y, sr, p.sr)\n",
    "    y = np.concatenate((np.zeros(512,),y),0)\n",
    "    new_filename = os.path.basename(audiofile)\n",
    "    new_filename = new_filename.replace('.wav','dfn_48k'+str(buffer_frame)+'.wav')\n",
    "    \n",
    "    Y = analysis_frame(y, normalize=True)\n",
    "    \n",
    "    Y_erb = (np.real(Y)**2 + np.imag(Y)**2) @ ERB_Matrix\n",
    "#     Y_erb = tf.sqrt(Y_erb)\n",
    "    Y_erb = 10*np.log10(Y_erb + 1e-10)\n",
    "    Y_erb_norm = erb_norm(Y_erb)\n",
    "    \n",
    "#     Y_spec_norm = as_real(Y[:,:p.nb_df])\n",
    "    Y_spec_norm = as_real(unit_norm(Y[:,:p.nb_df]))\n",
    "                            \n",
    "    buffer = np.zeros((buffer_frame,p.fft_size//2+1), dtype=complex)\n",
    "    df_buffer = np.zeros((1, p.df_order,p.fft_size//2+1,2))\n",
    "    buffer_erb = np.zeros((buffer_frame+2,32,1))\n",
    "    buffer_spec = np.zeros((buffer_frame+2,96,2))\n",
    "    print(Y.shape)\n",
    "    Z = np.zeros_like(Y, dtype=complex)\n",
    "    zg = np.zeros_like(Y, dtype=complex)\n",
    "    \n",
    "    last_df_alpha = np.zeros((1,1,1))\n",
    "    last_df_coeff = np.zeros((1,1,p.df_order,p.nb_df,2))\n",
    "    \n",
    "    state0, state10, state11, state20, state21 = np.zeros((1,256)), np.zeros((1,256)), np.zeros((1,256)), \\\n",
    "                                                 np.zeros((1,256)), np.zeros((1,256))\n",
    "    rstate0, rstate10, rstate11, rstate20, rstate21 = [],[],[],[],[]\n",
    "    ZG = np.zeros_like(Y, dtype=complex)\n",
    "    mask_record = np.zeros((1,1,p.fft_size//2+1,1))\n",
    "    \n",
    "    if not stateful:\n",
    "        for i in range(Y.shape[0]):\n",
    "            buffer = np.roll(buffer, -1, axis=0)\n",
    "            if i>0:\n",
    "                buffer[-1] = Y[i-1]\n",
    "\n",
    "            buffer_erb = np.roll(buffer_erb, -1, axis=0)\n",
    "            buffer_erb[-1] = np.reshape(Y_erb_norm[i],(p.nb_erb,1))\n",
    "\n",
    "            buffer_spec = np.roll(buffer_spec, -1, axis=0)\n",
    "            buffer_spec[-1] = Y_spec_norm[i]\n",
    "\n",
    "            \n",
    "            buffer_feature = [np.expand_dims(buffer_erb,0), np.expand_dims(buffer_spec,0),\n",
    "                         state0, [state10, state11], [state20, state21]]\n",
    "\n",
    "            rstate0.append(state0)\n",
    "            rstate10.append(state10)\n",
    "            rstate11.append(state11)\n",
    "            rstate20.append(state20)\n",
    "            rstate21.append(state21)\n",
    "\n",
    "            mask_out, df_alpha, df_coeff, state0, [state10, state11] ,[state20, state21] = model.predict(buffer_feature, verbose=0)\n",
    "\n",
    "            state0 = np.reshape(state0, (1,256))\n",
    "            state10 = np.reshape(state10, (1,256))\n",
    "            state11 = np.reshape(state11, (1,256))\n",
    "            state20 = np.reshape(state20, (1,256))\n",
    "            state21 = np.reshape(state21, (1,256))\n",
    "\n",
    "            buffer_real = np.concatenate((np.expand_dims(np.real(buffer),-1),np.expand_dims(np.imag(buffer),-1)),axis=-1)\n",
    "            buffer_spec_proc, mask = mask_operations(buffer_real, mask_out, training=False, return_mask=True, \n",
    "                                                     _min=0.0, cutoff=0.0) # mask gain\n",
    "\n",
    "            mask_record = np.concatenate((mask_record,mask),0)\n",
    "\n",
    "            df_buffer = np.roll(df_buffer, -1, axis=1)\n",
    "            df_buffer[:,-1] = buffer_spec_proc\n",
    "\n",
    "            df_alpha = tf.cast(df_alpha, tf.float64)\n",
    "            enhanced = df_operations(df_buffer, last_df_coeff, last_df_alpha, training=False) # deep filter\n",
    "    #         enhanced = buffer_spec_proc\n",
    "\n",
    "            enhanced = np.squeeze(enhanced)\n",
    "            enhanced_complex = np.squeeze(enhanced[..., 0]+1j*enhanced[..., 1]) \n",
    "\n",
    "            lim = 0.0\n",
    "            if lim>0: enhanced_complex = (1-lim) * enhanced_complex + lim * buffer \n",
    "\n",
    "            last_df_alpha = df_alpha\n",
    "            last_df_coeff = df_coeff\n",
    "\n",
    "            if buffer_frame==1:\n",
    "                Z[i,:] = enhanced_complex \n",
    "            else:\n",
    "                Z[i,:] = enhanced_complex[-1] \n",
    "    else:\n",
    "        buffer_feature = [Y_erb_norm, Y_spec_norm]\n",
    "        mask_out, df_alpha, df_coeff = model.predict(buffer_feature, verbose=0)\n",
    "        \n",
    "        Y = np.concatenate((np.zeros(2,513,1),Y),axis=0)\n",
    "        buffer_real = np.concatenate((np.expand_dims(np.real(Y),-1),np.expand_dims(np.imag(Y),-1)),axis=-1)\n",
    "        buffer_spec_proc, mask_record = mask_operations(buffer_real, mask_out, training=False, return_mask=True, \n",
    "                                                 _min=0.0, cutoff=0.0) # mask gain\n",
    "\n",
    "\n",
    "        df_alpha = tf.cast(df_alpha, tf.float64)\n",
    "        enhanced = df_operations(buffer_spec_proc, df_coeff, df_alpha, training=False) # deep filter\n",
    "\n",
    "        enhanced = np.squeeze(enhanced)\n",
    "        enhanced_complex = np.squeeze(enhanced[..., 0]+1j*enhanced[..., 1]) \n",
    "            \n",
    "    z = synthesis_frame(Z, windowing=True, normalize=True)\n",
    "\n",
    "    if not os.path.exists(process_path): os.makedirs(process_path)\n",
    "    y = np.reshape(y,(y.shape[0],1))\n",
    "    z = np.reshape(z,(z.shape[0],1))\n",
    "    if y.shape[0]>z.shape[0]:\n",
    "        z = np.concatenate((np.zeros((y.shape[0]-z.shape[0],1)),z),axis=0)\n",
    "    if y.shape[0]<z.shape[0]:\n",
    "        y = np.concatenate((np.zeros((z.shape[0]-y.shape[0],1)),y),axis=0)\n",
    "    \n",
    "    write(os.path.join(process_path,new_filename), np.concatenate((y,z),axis=1), p.sr)\n",
    "    \n",
    "    plt.figure(figsize=(10,5))\n",
    "    plt.plot(y, color = 'b')\n",
    "    plt.plot(z, color = 'r')\n",
    "    plt.legend(['unproc','proc'])\n",
    "    plt.show()\n",
    "    \n",
    "    plt.figure(figsize=(10,5))\n",
    "    plt.imshow(np.flipud(np.squeeze(mask_record).T), vmin=0.0, vmax=1.0)\n",
    "    plt.colorbar()\n",
    "    plt.show()\n",
    "    \n",
    "    plt.figure(figsize=(10,5))\n",
    "    plt.plot(np.mean(np.squeeze(mask_record),axis=1))\n",
    "    plt.hlines(0.02,xmin=0,xmax=mask_record.shape[0],colors='r')\n",
    "    plt.show()\n",
    "    \n",
    "    if stateful:\n",
    "        return Y.shape[0], \\\n",
    "               Y_erb_norm, \\\n",
    "               Y_spec_norm\n",
    "    else:\n",
    "        return Y.shape[0], \\\n",
    "               Y_erb_norm, \\\n",
    "               Y_spec_norm, \\\n",
    "               np.concatenate(rstate0, 0),\\\n",
    "               np.concatenate(rstate10, 0),\\\n",
    "               np.concatenate(rstate11, 0) ,\\\n",
    "               np.concatenate(rstate20, 0),\\\n",
    "               np.concatenate(rstate21, 0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "id": "e9ed5389",
   "metadata": {},
   "outputs": [],
   "source": [
    "import time\n",
    "buffer_frame = 1"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "id": "8d2622bd",
   "metadata": {},
   "outputs": [],
   "source": [
    "process_path = '/home/myhsueh/DeepFilterNet/wav/proc/'\n",
    "\n",
    "audiofile = [\n",
    "        '/home/myhsueh/ABR/wav/Respeaker_mic2.wav',\n",
    "        '/home/myhsueh/ABR/wav/Respeaker_mic2_moving.wav', \n",
    "#               '/home/myhsueh/DeepFilterNet/Keras_DeepFilterNet2_1024_48k/8368P_tflite/aec_20_out.wav',\n",
    "#               '/home/myhsueh/DeepFilterNet/wav/orig/cafe_test2.wav',\n",
    "#               '/home/myhsueh/DeepFilterNet/wav/orig/Cafe_noise_1min43s.wav',\n",
    "#               '/home/myhsueh/DeepFilterNet/wav/orig/recordwBGN.wav',\n",
    "#               '/home/myhsueh/DeepFilterNet/wav/orig/OURS_all_off.wav',\n",
    "#               '/home/myhsueh/DeepFilterNet/wav/orig/OURS_all_off_2.wav',\n",
    "    \n",
    "#               '/home/myhsueh/DeepFilterNet/wav/orig/DNS_sample3_01.wav',\n",
    "#               '/home/myhsueh/DeepFilterNet/wav/orig/DNS_sample1_01.wav',\n",
    "#               '/home/myhsueh/DeepFilterNet/wav/orig/DNS_sample2_01.wav',\n",
    "#               '/home/myhsueh/DeepFilterNet/wav/orig/DNS_sample6_01.wav',\n",
    "#               '/home/myhsueh/DeepFilterNet/wav/orig/DNS_sample7_01.wav',\n",
    "#               '/home/myhsueh/DeepFilterNet/wav/orig/DNS_sample9_01.wav',\n",
    "#               '/home/myhsueh/DeepFilterNet/wav/orig/DNS_sample11_01.wav',\n",
    "    \n",
    "#               '/home/myhsueh/DeepFilterNet/wav/orig/DNS_testset/ms_realrec_english_usbmicrophone_APPND34J4EDS3_typing_near_fileid_5.wav',\n",
    "#               '/home/myhsueh/DeepFilterNet/wav/orig/DNS_testset/ms_realrec_english_speaker_A33BAVCUPSMTWJ_Typing_near_fileid_11.wav',\n",
    "#               '/home/myhsueh/DeepFilterNet/wav/orig/DNS_testset/ms_realrec_english_openspeaker_ASEW6NZHLI41K_typing_far_fileid_9.wav',\n",
    "#               '/home/myhsueh/DeepFilterNet/wav/orig/DNS_testset/ms_realrec_english_openspeaker_ASEW6NZHLI41K_openingchipspacket_near_fileid_1.wav',\n",
    "#               '/home/myhsueh/DeepFilterNet/wav/orig/DNS_testset/ms_realrec_english_openspeaker_A338ZXK723N9UH_dogbarking_near_fileid_si2192.wav',\n",
    "#               '/home/myhsueh/DeepFilterNet/wav/orig/DNS_testset/ms_realrec_english_openspeaker_A338ZXK723N9UH_babycrying_far_fileid_si1687.wav',\n",
    "#               '/home/myhsueh/DeepFilterNet/wav/orig/DNS_testset/ms_realrec_english_openspeaker_A338ZXK723N9UH_ArConditioner_near_fileid_si2313.wav',\n",
    "#               '/home/myhsueh/DeepFilterNet/wav/orig/DNS_testset/ms_realrec_english_laptopmic_A9AP3ED0K8LS8_heavybreathing_near_fileid_6.wav',\n",
    "#               '/home/myhsueh/DeepFilterNet/wav/orig/DNS_testset/ms_realrec_emotional_Desktopstandmic_AHLK9SWDJHBBZ_clatternoise_near_crying_fileid_8.wav',\n",
    "               \n",
    "        \n",
    "            ]\n",
    "import glob"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 29,
   "id": "380b2a11",
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "<ipython-input-26-d957bbf90d18>:3: FutureWarning: Pass orig_sr=16000, target_sr=48000 as keyword args. From version 0.10 passing these as positional arguments will result in an error\n",
      "  if not sr==p.sr: y = librosa.resample(y, sr, p.sr)\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "(22237, 513)\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 41276 characters]\n",
      "text/plain": [
       "<Figure size 720x360 with 1 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    },
    {
     "data": {
      "image/png": "[encoded data, 17212 characters]\n",
      "text/plain": [
       "<Figure size 720x360 with 2 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    },
    {
     "data": {
      "image/png": "[encoded data, 35120 characters]\n",
      "text/plain": [
       "<Figure size 720x360 with 1 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "0.09346515960700222\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "<ipython-input-26-d957bbf90d18>:3: FutureWarning: Pass orig_sr=16000, target_sr=48000 as keyword args. From version 0.10 passing these as positional arguments will result in an error\n",
      "  if not sr==p.sr: y = librosa.resample(y, sr, p.sr)\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "(24427, 513)\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 43604 characters]\n",
      "text/plain": [
       "<Figure size 720x360 with 1 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    },
    {
     "data": {
      "image/png": "[encoded data, 17500 characters]\n",
      "text/plain": [
       "<Figure size 720x360 with 2 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    },
    {
     "data": {
      "image/png": "[encoded data, 33716 characters]\n",
      "text/plain": [
       "<Figure size 720x360 with 1 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "0.1024593671833141\n"
     ]
    }
   ],
   "source": [
    "import time\n",
    "model.reset_states()  \n",
    "    \n",
    "r_erb_list, r_spec_list = [],[]\n",
    "r_state_list, r_state10_list, r_state11_list, r_state20_list, r_state21_list = [],[],[],[],[]\n",
    "buffer_frame = 1\n",
    "for file in audiofile:\n",
    "    start = time.time()\n",
    "    if stateful:\n",
    "        num_frame, r_erb, r_spec = process(file)\n",
    "    else:\n",
    "        num_frame, r_erb, r_spec, r_state, r_state10, r_state11, r_state20, r_state21 = process(file)\n",
    "    end = time.time()\n",
    "    print((end-start)/num_frame)\n",
    "    \n",
    "    r_erb_list.append(r_erb)\n",
    "    r_spec_list.append(r_spec)\n",
    "    if not stateful:\n",
    "        r_state_list.append(r_state)\n",
    "        r_state10_list.append(r_state10)\n",
    "        r_state11_list.append(r_state11)\n",
    "        r_state20_list.append(r_state20)\n",
    "        r_state21_list.append(r_state21)\n",
    "\n",
    "# if not os.path.isdir('./quantize_npy'): os.makedirs('./quantize_npy')\n",
    "# np.save('./quantize_npy/whole_erb.npy', np.concatenate(r_erb_list,0))\n",
    "# np.save('./quantize_npy/whole_spec.npy',np.concatenate(r_spec_list,0))\n",
    "# np.save('./quantize_npy/whole_state.npy',np.concatenate(r_state_list,0))\n",
    "# np.save('./quantize_npy/whole_state10.npy',np.concatenate(r_state10_list,0))\n",
    "# np.save('./quantize_npy/whole_state11.npy',np.concatenate(r_state11_list,0))\n",
    "# np.save('./quantize_npy/whole_state20.npy',np.concatenate(r_state20_list,0))\n",
    "# np.save('./quantize_npy/whole_state21.npy',np.concatenate(r_state21_list,0))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 30,
   "id": "51b9bad0",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "INFO:numba.cuda.cudadrv.driver:init\n",
      "ERROR:numba.cuda.cudadrv.driver:Call to cuInit results in CUDA_ERROR_NO_DEVICE\n"
     ]
    },
    {
     "ename": "CudaSupportError",
     "evalue": "Error at driver init: \n[100] Call to cuInit results in CUDA_ERROR_NO_DEVICE:",
     "output_type": "error",
     "traceback": [
      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[0;31mCudaAPIError\u001b[0m                              Traceback (most recent call last)",
      "\u001b[0;32m/opt/anaconda3/lib/python3.8/site-packages/numba/cuda/cudadrv/driver.py\u001b[0m in \u001b[0;36minitialize\u001b[0;34m(self)\u001b[0m\n\u001b[1;32m    236\u001b[0m             \u001b[0m_logger\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0minfo\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m'init'\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 237\u001b[0;31m             \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mcuInit\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;36m0\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m    238\u001b[0m         \u001b[0;32mexcept\u001b[0m \u001b[0mCudaAPIError\u001b[0m \u001b[0;32mas\u001b[0m \u001b[0me\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m/opt/anaconda3/lib/python3.8/site-packages/numba/cuda/cudadrv/driver.py\u001b[0m in \u001b[0;36msafe_cuda_api_call\u001b[0;34m(*args)\u001b[0m\n\u001b[1;32m    299\u001b[0m             \u001b[0mretcode\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mlibfn\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m*\u001b[0m\u001b[0margs\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 300\u001b[0;31m             \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m_check_error\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mfname\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mretcode\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m    301\u001b[0m         \u001b[0;32mreturn\u001b[0m \u001b[0msafe_cuda_api_call\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m/opt/anaconda3/lib/python3.8/site-packages/numba/cuda/cudadrv/driver.py\u001b[0m in \u001b[0;36m_check_error\u001b[0;34m(self, fname, retcode)\u001b[0m\n\u001b[1;32m    334\u001b[0m                     \u001b[0;32mraise\u001b[0m \u001b[0mCudaDriverError\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m\"CUDA initialized before forking\"\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 335\u001b[0;31m             \u001b[0;32mraise\u001b[0m \u001b[0mCudaAPIError\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mretcode\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mmsg\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m    336\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;31mCudaAPIError\u001b[0m: [100] Call to cuInit results in CUDA_ERROR_NO_DEVICE",
      "\nDuring handling of the above exception, another exception occurred:\n",
      "\u001b[0;31mCudaSupportError\u001b[0m                          Traceback (most recent call last)",
      "\u001b[0;32m<ipython-input-30-d1b10d799ccb>\u001b[0m in \u001b[0;36m<module>\u001b[0;34m\u001b[0m\n\u001b[1;32m      1\u001b[0m \u001b[0;32mfrom\u001b[0m \u001b[0mnumba\u001b[0m \u001b[0;32mimport\u001b[0m \u001b[0mcuda\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m----> 2\u001b[0;31m \u001b[0mdevice\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0mcuda\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mget_current_device\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m      3\u001b[0m \u001b[0mdevice\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mreset\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m/opt/anaconda3/lib/python3.8/site-packages/numba/cuda/api.py\u001b[0m in \u001b[0;36mget_current_device\u001b[0;34m()\u001b[0m\n\u001b[1;32m    456\u001b[0m \u001b[0;32mdef\u001b[0m \u001b[0mget_current_device\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    457\u001b[0m     \u001b[0;34m\"Get current device associated with the current thread\"\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 458\u001b[0;31m     \u001b[0;32mreturn\u001b[0m \u001b[0mcurrent_context\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mdevice\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m    459\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    460\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m/opt/anaconda3/lib/python3.8/site-packages/numba/cuda/cudadrv/devices.py\u001b[0m in \u001b[0;36mget_context\u001b[0;34m(devnum)\u001b[0m\n\u001b[1;32m    210\u001b[0m     \u001b[0;32mreturn\u001b[0m \u001b[0mthe\u001b[0m \u001b[0mCUDA\u001b[0m \u001b[0mcontext\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    211\u001b[0m     \"\"\"\n\u001b[0;32m--> 212\u001b[0;31m     \u001b[0;32mreturn\u001b[0m \u001b[0m_runtime\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mget_or_create_context\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mdevnum\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m    213\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    214\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m/opt/anaconda3/lib/python3.8/site-packages/numba/cuda/cudadrv/devices.py\u001b[0m in \u001b[0;36mget_or_create_context\u001b[0;34m(self, devnum)\u001b[0m\n\u001b[1;32m    136\u001b[0m             \u001b[0mattached_ctx\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m_get_attached_context\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    137\u001b[0m             \u001b[0;32mif\u001b[0m \u001b[0mattached_ctx\u001b[0m \u001b[0;32mis\u001b[0m \u001b[0;32mNone\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 138\u001b[0;31m                 \u001b[0;32mreturn\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m_get_or_create_context_uncached\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mdevnum\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m    139\u001b[0m             \u001b[0;32melse\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    140\u001b[0m                 \u001b[0;32mreturn\u001b[0m \u001b[0mattached_ctx\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m/opt/anaconda3/lib/python3.8/site-packages/numba/cuda/cudadrv/devices.py\u001b[0m in \u001b[0;36m_get_or_create_context_uncached\u001b[0;34m(self, devnum)\u001b[0m\n\u001b[1;32m    149\u001b[0m             \u001b[0;31m# Try to get the active context in the CUDA stack or\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    150\u001b[0m             \u001b[0;31m# activate GPU-0 with the primary context\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 151\u001b[0;31m             \u001b[0;32mwith\u001b[0m \u001b[0mdriver\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mget_active_context\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m \u001b[0;32mas\u001b[0m \u001b[0mac\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m    152\u001b[0m                 \u001b[0;32mif\u001b[0m \u001b[0;32mnot\u001b[0m \u001b[0mac\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    153\u001b[0m                     \u001b[0;32mreturn\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m_activate_context_for\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;36m0\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m/opt/anaconda3/lib/python3.8/site-packages/numba/cuda/cudadrv/driver.py\u001b[0m in \u001b[0;36m__enter__\u001b[0;34m(self)\u001b[0m\n\u001b[1;32m    391\u001b[0m         \u001b[0;32melse\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    392\u001b[0m             \u001b[0mhctx\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mdrvapi\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mcu_context\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;36m0\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 393\u001b[0;31m             \u001b[0mdriver\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mcuCtxGetCurrent\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mbyref\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mhctx\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m    394\u001b[0m             \u001b[0mhctx\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mhctx\u001b[0m \u001b[0;32mif\u001b[0m \u001b[0mhctx\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mvalue\u001b[0m \u001b[0;32melse\u001b[0m \u001b[0;32mNone\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    395\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m/opt/anaconda3/lib/python3.8/site-packages/numba/cuda/cudadrv/driver.py\u001b[0m in \u001b[0;36m__getattr__\u001b[0;34m(self, fname)\u001b[0m\n\u001b[1;32m    278\u001b[0m         \u001b[0;31m# Initialize driver\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    279\u001b[0m         \u001b[0;32mif\u001b[0m \u001b[0;32mnot\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mis_initialized\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 280\u001b[0;31m             \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0minitialize\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m    281\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    282\u001b[0m         \u001b[0;32mif\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0minitialization_error\u001b[0m \u001b[0;32mis\u001b[0m \u001b[0;32mnot\u001b[0m \u001b[0;32mNone\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m/opt/anaconda3/lib/python3.8/site-packages/numba/cuda/cudadrv/driver.py\u001b[0m in \u001b[0;36minitialize\u001b[0;34m(self)\u001b[0m\n\u001b[1;32m    238\u001b[0m         \u001b[0;32mexcept\u001b[0m \u001b[0mCudaAPIError\u001b[0m \u001b[0;32mas\u001b[0m \u001b[0me\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    239\u001b[0m             \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0minitialization_error\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0me\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 240\u001b[0;31m             \u001b[0;32mraise\u001b[0m \u001b[0mCudaSupportError\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m\"Error at driver init: \\n%s:\"\u001b[0m \u001b[0;34m%\u001b[0m \u001b[0me\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m    241\u001b[0m         \u001b[0;32melse\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    242\u001b[0m             \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mpid\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0m_getpid\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;31mCudaSupportError\u001b[0m: Error at driver init: \n[100] Call to cuInit results in CUDA_ERROR_NO_DEVICE:"
     ]
    }
   ],
   "source": [
    "from numba import cuda\n",
    "device=cuda.get_current_device()\n",
    "device.reset()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7fac21b7",
   "metadata": {},
   "outputs": [],
   "source": [
    "converter = tf.lite.TFLiteConverter.from_keras_model(model)\n",
    "\n",
    "converter.optimizations = [tf.lite.Optimize.DEFAULT]\n",
    "tflite_model_quant = converter.convert()\n",
    "\n",
    "with open('model.tflite', 'wb') as f:\n",
    "    f.write(tflite_model_quant)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "152abe1a",
   "metadata": {},
   "source": [
    "# Quant"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2a5d9ff5",
   "metadata": {},
   "outputs": [],
   "source": [
    "quant_erb = np.load('./quantize_npy/whole_erb_half.npy')\n",
    "quant_spec = np.load('./quantize_npy/whole_spec_half.npy')\n",
    "quant_state0 = np.load('./quantize_npy/whole_state_half.npy')\n",
    "quant_state10 = np.load('./quantize_npy/whole_state10_half.npy')\n",
    "quant_state11 = np.load('./quantize_npy/whole_state11_half.npy')\n",
    "quant_state20 = np.load('./quantize_npy/whole_state20_half.npy')\n",
    "quant_state21 = np.load('./quantize_npy/whole_state21_half.npy')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a2358adf",
   "metadata": {},
   "outputs": [],
   "source": [
    "quant_erb = np.concatenate((np.zeros((2,32,1)),quant_erb),axis=0)\n",
    "quant_spec = np.concatenate((np.zeros((2,96,2)),quant_spec),axis=0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e3fa8927",
   "metadata": {},
   "outputs": [],
   "source": [
    "range_list = [i for i in range(quant_state21.shape[0])]\n",
    "random.shuffle(range_list)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "475a534c",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "26fb72aa",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "def representative_data_gen():\n",
    "    for j in range_list:\n",
    "        yield [\n",
    "               quant_spec[j:j+3].reshape(1,3,96,2).astype(np.float32), \n",
    "               quant_state11[j].reshape(1,256).astype(np.float32), \n",
    "               quant_state10[j].reshape(1,256).astype(np.float32), \n",
    "               quant_state21[j].reshape(1,256).astype(np.float32), \n",
    "               quant_state20[j].reshape(1,256).astype(np.float32),\n",
    "               quant_erb[j:j+3].reshape(1,3,32,1).astype(np.float32), \n",
    "               quant_state0[j].reshape(1,256).astype(np.float32)\n",
    "        ]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "91b9b748",
   "metadata": {},
   "outputs": [],
   "source": [
    "converter = tf.lite.TFLiteConverter.from_keras_model(model)\n",
    "\n",
    "converter.optimizations = [tf.lite.Optimize.DEFAULT]\n",
    "converter.representative_dataset = representative_data_gen\n",
    "# Ensure that if any ops can't be quantized, the converter throws an error\n",
    "converter.target_spec.supported_ops = [tf.lite.OpsSet.EXPERIMENTAL_TFLITE_BUILTINS_ACTIVATIONS_INT16_WEIGHTS_INT8] \n",
    "# converter.target_spec.supported_ops = [tf.lite.OpsSet.TFLITE_BUILTINS_INT8]\n",
    "# Set the input and output tensors to uint8 (APIs added in r2.3)\n",
    "converter.inference_input_type = tf.int16\n",
    "converter.inference_output_type = tf.int16\n",
    "\n",
    "tflite_model_quant = converter.convert()\n",
    "with open('INT16_model.tflite', 'wb') as f:\n",
    "    f.write(tflite_model_quant)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ecd1ad07",
   "metadata": {},
   "outputs": [],
   "source": [
    "converter = tf.lite.TFLiteConverter.from_keras_model(model)\n",
    "\n",
    "converter.optimizations = [tf.lite.Optimize.DEFAULT]\n",
    "\n",
    "tflite_model_quant = converter.convert()\n",
    "with open('FP32_model.tflite', 'wb') as f:\n",
    "    f.write(tflite_model_quant)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3afa2904",
   "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.8.10"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
