STARN_model.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "import tensorflow as tf\n",
    "from tensorflow import keras\n",
    "from keras.layers import (\n",
    "    Conv2D,\n",
    "    Conv3D,\n",
    "    GlobalAveragePooling3D,\n",
    "    GlobalMaxPooling3D,\n",
    "    BatchNormalization, \n",
    "    Activation, \n",
    "    add,\n",
    "    Reshape,\n",
    "    multiply,\n",
    "    Lambda,\n",
    "    concatenate,\n",
    "    Input,\n",
    "    Permute,\n",
    "    Dropout\n",
    ")\n",
    "from keras.models import Model\n",
    "import keras.backend as K\n",
    "from keras import layers, regularizers\n",
    "\n",
    "physical_devices = tf.config.list_physical_devices('GPU')\n",
    "\n",
    "tf.config.experimental.set_visible_devices(devices=physical_devices[0], device_type='GPU')\n",
    "tf.config.experimental.set_memory_growth(physical_devices[0], enable=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "keras.backend.set_image_data_format('channels_last')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Attention Block"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "# channel last\n",
    "# for 3D model\n",
    "class ChannelAttention(layers.Layer):\n",
    "    \"\"\"\n",
    "    inputs: (None,height,width,time,channel)\n",
    "    channel_attention\n",
    "    \"\"\"\n",
    "    def __init__(self,channels,ratio=16,**kwargs):\n",
    "        super(ChannelAttention,self).__init__(**kwargs)\n",
    "        self.avg = GlobalAveragePooling3D()\n",
    "        self.max = GlobalMaxPooling3D()\n",
    "        self.channels = channels\n",
    "        self.ratio = ratio\n",
    "        self.conv1 = Conv3D(self.channels//self.ratio,kernel_size=1,strides=1,padding='same',\n",
    "                            activation='relu',use_bias=True,bias_initializer='zeros',\n",
    "                            kernel_initializer='he_normal',name = 'CA_Conv1')\n",
    "        self.conv2 = Conv3D(self.channels,kernel_size=1,strides=1,padding='same',use_bias=True,\n",
    "                            bias_initializer='zeros',kernel_initializer='he_normal',name = 'CA_Conv2')\n",
    "\n",
    "    def call(self,inputs):\n",
    "        #channel = inputs.get_shape().as_list()[-1]\n",
    "        avg = self.avg(inputs)\n",
    "        max = self.max(inputs)\n",
    "        avg = Reshape((1,1,1,avg.shape[1]))(avg) # shape (None,1,1,1,channel)\n",
    "        max = Reshape((1,1,1,max.shape[1]))(max)\n",
    "        avg_out = self.conv2(self.conv1(avg))\n",
    "        max_out = self.conv2(self.conv1(max))\n",
    "        out = add([avg_out,max_out])\n",
    "        out = Activation('sigmoid')(out)\n",
    "        return multiply([inputs,out]) #out #\n",
    "    \n",
    "    def get_config(self):\n",
    "        config = super().get_config().copy()\n",
    "        config.update({\n",
    "            'channels': self.channels,\n",
    "            'ratio': self.ratio\n",
    "        })\n",
    "        return config\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "class SpatialAttention(layers.Layer):\n",
    "    \"\"\"\n",
    "    spatial_attention\n",
    "    kernel_size: kernel size for spatial domain\n",
    "    \"\"\"\n",
    "    def __init__(self,kernel_size = 5,**kwargs): #之前为7\n",
    "        super(SpatialAttention,self).__init__(**kwargs)\n",
    "        self.kernel_size = kernel_size\n",
    "        #卷积核数应为(5,5,1) 在时间维度上是1  \n",
    "        self.conv1 = Conv3D(filters= 1,kernel_size=(self.kernel_size,self.kernel_size,1),strides=1,padding='same',\n",
    "                            activation='sigmoid',kernel_initializer='he_normal',use_bias=False,name = 'SA_Conv1')\n",
    "    \n",
    "    def call(self,inputs):\n",
    "        #输出shape:(None,H,W,T,1) axis = -1\n",
    "        #输出应该为:(None,H,W,1,1) axis = [3,4]\n",
    "        avg_out = Lambda(lambda x: K.mean(x, axis = [3,4], keepdims=True))(inputs) \n",
    "        max_out = Lambda(lambda x: K.max(x, axis = [3,4], keepdims=True))(inputs)\n",
    "        concat = concatenate([avg_out,max_out],axis = -1)\n",
    "        out = self.conv1(concat)\n",
    "        return multiply([inputs,out]) #out #\n",
    "    \n",
    "    def get_config(self):\n",
    "        config = super().get_config().copy()\n",
    "        config.update({\n",
    "            'kernel_size': self.kernel_size\n",
    "        })\n",
    "        return config\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "class TemporalAttention(layers.Layer):\n",
    "    \"\"\"\n",
    "    channels: 输入的Conv中的通道数\n",
    "    temporal_attention\n",
    "    注意Time是第几维\n",
    "    input shape: (None,H,W,T,C)\n",
    "    channels:输入的time维度的通道数\n",
    "    \"\"\"\n",
    "    def __init__(self,channels,**kwargs):\n",
    "        super(TemporalAttention,self).__init__(**kwargs)\n",
    "        self.avg = GlobalAveragePooling3D()\n",
    "        self.max = GlobalMaxPooling3D()\n",
    "        self.channels = channels\n",
    "        self.conv1 = Conv3D(self.channels,kernel_size=1,strides=1,padding='same',\n",
    "                            activation='relu',use_bias=True,bias_initializer='zeros',\n",
    "                            kernel_initializer='he_normal',name = 'TA_Conv1')\n",
    "        self.conv2 = Conv3D(self.channels,kernel_size=1,strides=1,padding='same',use_bias=True,\n",
    "                            bias_initializer='zeros',kernel_initializer='he_normal',name = 'TA_Conv2')\n",
    "\n",
    "    def call(self,inputs):\n",
    "        x = Permute((1,2,4,3))(inputs)\n",
    "        avg = self.avg(x)\n",
    "        max = self.max(x)\n",
    "        avg = Reshape((1,1,1,avg.shape[1]))(avg) # shape (None,1,1,1,time)\n",
    "        max = Reshape((1,1,1,max.shape[1]))(max)\n",
    "        avg_out = self.conv2(self.conv1(avg)) # 这里卷积核为(1,1,1)所以time还是最后一个维度\n",
    "        max_out = self.conv2(self.conv1(max)) # 输出为 （None,1,1,1,64)\n",
    "        avg_out = Permute((1,2,4,3))(avg_out) # change shape to (None,1,1,time,1)\n",
    "        max_out = Permute((1,2,4,3))(max_out)\n",
    "        out = add([avg_out,max_out])\n",
    "        out = Activation('sigmoid')(out)\n",
    "        return multiply([inputs,out]) #out #\n",
    "\n",
    "    def get_config(self):\n",
    "        config = super().get_config().copy()\n",
    "        config.update({\n",
    "            'channels': self.channels\n",
    "        })\n",
    "        return config\n",
    "        "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "class RDAB(Model):\n",
    "    \"\"\"\n",
    "    \n",
    "    \"\"\"\n",
    "    def __init__(self,channels,t_dims,**kwargs):\n",
    "        super(RDAB,self).__init__(**kwargs)\n",
    "        \n",
    "        self.channels = channels\n",
    "        self.t_dims = t_dims\n",
    "        self.c1 = Conv3D(filters=self.channels,kernel_initializer='he_normal',kernel_size=(3,3,3),activation='relu',padding='same',name = 'RDAB_Conv1')\n",
    "        self.c2 = Conv3D(filters=self.channels,kernel_initializer='he_normal',kernel_size=(3,3,3),activation='relu',padding='same',name = 'RDAB_Conv2')\n",
    "        self.c3 = Conv3D(filters=self.channels,kernel_initializer='he_normal',kernel_size=(3,3,3),activation='relu',padding='same',name = 'RDAB_Conv3')\n",
    "        self.c4 = Conv3D(filters=self.channels,kernel_initializer='he_normal',kernel_size=(3,3,3),activation='relu',padding='same',name = 'RDAB_Conv4')\n",
    "        self.c5 = Conv3D(filters=self.channels,kernel_initializer='he_normal',kernel_size=(3,3,3),activation='relu',padding='same',name = 'RDAB_Conv5')\n",
    "        self.c6 = Conv3D(filters=self.channels,kernel_initializer='he_normal',kernel_size=(3,3,3),activation='relu',padding='same',name = 'RDAB_Conv6')\n",
    "        self.c = Conv3D(filters=64,kernel_initializer='he_normal',kernel_size=(1,1,1),padding='same',name = 'RDAB_Conv7')\n",
    "        #ChannelAttention中的channel要和前面最后一个的输出对应\n",
    "        self.ca = ChannelAttention(self.channels)\n",
    "        self.ta = TemporalAttention(self.t_dims)\n",
    "        self.sa = SpatialAttention()\n",
    "\n",
    "\n",
    "    def call(self,inputs):\n",
    "        x1 = self.c1(inputs)\n",
    "        #在channel维进行concatenate\n",
    "        y1 = tf.concat([inputs,x1],axis=-1)\n",
    "\n",
    "        x2 = self.c2(y1)\n",
    "        y2 = tf.concat([inputs,x1,x2],axis=-1)\n",
    "\n",
    "        x3 = self.c3(y2)\n",
    "        y3 = tf.concat([inputs,x1,x2,x3],axis=-1)\n",
    "\n",
    "        x4 = self.c4(y3)\n",
    "        y4 = tf.concat([inputs,x1,x2,x3,x4],axis=-1)\n",
    "        \n",
    "        x5 = self.c5(y4)\n",
    "        y5 = tf.concat([inputs,x1,x2,x3,x4,x5],axis=-1)\n",
    "\n",
    "        x6 = self.c6(y5)\n",
    "        y6 = tf.concat([inputs,x1,x2,x3,x4,x5,x6],axis=-1)\n",
    "\n",
    "        y = self.c(y6)\n",
    "\n",
    "        #注意力是串行的\n",
    "        ca_out = self.ca(y)\n",
    "        sa_out = self.sa(ca_out)\n",
    "        ta_out = self.ta(sa_out)\n",
    "\n",
    "        out = add([ta_out,inputs])\n",
    "        return out\n",
    "\n",
    "  \n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### STARN"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [],
   "source": [
    "class STARN(Model):\n",
    "    \"\"\"\n",
    "    endocer-decoder RDN with CBAM\n",
    "    \"\"\"\n",
    "    def __init__(self,input_shape,num_G=64,**kwargs):\n",
    "        super(STARN,self).__init__(**kwargs)\n",
    "        self.num_G = num_G\n",
    "        self.c_dims = input_shape[3]\n",
    "        self.t_dims = input_shape[2]\n",
    "        #self.channels = channels\n",
    "        self.input_layer = Input(input_shape)\n",
    "        self.conv1 = Conv3D(filters=64,kernel_size=(3,3,3),kernel_regularizer = regularizers.l2(5e-4),padding='same',name = 'STARN_Conv1')\n",
    "        #,kernel_initializer='he_normal'\n",
    "        # encoder\n",
    "        self.Econv1 = Conv3D(filters=64,kernel_size=(3,3,3),padding='same',name = 'STARN_Econv1')\n",
    "        self.Ebn1 = BatchNormalization(momentum=0.8,name = 'STARN_Ebn1')\n",
    "        self.Econv2 = Conv3D(filters=64,kernel_size=(3,3,3),padding='same',name = 'STARN_Econv2')\n",
    "        self.Ebn2 = BatchNormalization(momentum=0.8,name = 'STARN_Ebn2')\n",
    "        self.Econv3 = Conv3D(filters=64,kernel_size=(3,3,3),padding='same',name = 'STARN_Econv3')\n",
    "        self.Ebn3 = BatchNormalization(momentum=0.8,name = 'STARN_Ebn3')\n",
    "        self.Econv4 = Conv3D(filters=64,kernel_size=(3,3,3),padding='same',name = 'STARN_Econv4')\n",
    "        self.Ebn4 = BatchNormalization(momentum=0.8,name = 'STARN_Ebn4')\n",
    "        \n",
    "        #decoder\n",
    "        self.Dconv1 = Conv3D(filters=64,kernel_size=(3,3,3),padding='same',name = 'STARN_Dconv1')\n",
    "        self.Dbn1 = BatchNormalization(momentum=0.8,name = 'STARN_Dbn1')\n",
    "        self.Dconv2 = Conv3D(filters=64,kernel_size=(3,3,3),padding='same',name = 'STARN_Dconv2')\n",
    "        self.Dbn2 = BatchNormalization(momentum=0.8,name = 'STARN_Dbn2')\n",
    "        self.Dconv3 = Conv3D(filters=64,kernel_size=(3,3,3),padding='same',name = 'STARN_Dconv3')\n",
    "        self.Dbn3 = BatchNormalization(momentum=0.8,name = 'STARN_Dbn3')\n",
    "        self.Dconv4 = Conv3D(filters=64,kernel_size=(3,3,3),padding='same',name = 'STARN_Dconv4')\n",
    "        self.Dbn4 = BatchNormalization(momentum=0.8,name = 'STARN_Dbn4')\n",
    "\n",
    "        #RDN_CBAM Block\n",
    "        self.Block1 = RDAB(self.num_G,self.t_dims)\n",
    "        self.Block2 = RDAB(self.num_G,self.t_dims)\n",
    "        self.Block3 = RDAB(self.num_G,self.t_dims)\n",
    "        self.Block4 = RDAB(self.num_G,self.t_dims)\n",
    "        #self.Block5 = RDAB(self.num_G,self.t_dims)\n",
    "        #self.Block6 = RDAB(self.num_G,self.t_dims)\n",
    "        \n",
    "        #activation sigmoid or linear\n",
    "        self.conv2 = Conv3D(filters=1,kernel_size=(3,3,3),padding='same',name = 'STARN_Conv2')\n",
    "        self.conv3 = Conv2D(filters=1,kernel_size=(2,2),padding='same',name = 'STARN_Conv3')\n",
    "        self.ta = TemporalAttention(self.t_dims)\n",
    "        self.bn1 = BatchNormalization(momentum=0.8,name = 'STARN_Bn1')\n",
    "        #self.drop1 = Dropout(0.2)\n",
    "        #self.drop2 = Dropout(0.2)\n",
    "        #self.drop3 = Dropout(0.2)\n",
    "        self.out = self.call(self.input_layer)\n",
    "        super(STARN,self).__init__(\n",
    "            inputs = self.input_layer,\n",
    "            outputs = self.out,\n",
    "            **kwargs\n",
    "        )\n",
    "\n",
    "    def call(self,inputs,training = False):\n",
    "        x = self.conv1(inputs)\n",
    "        x = Activation(activation='relu')(x)\n",
    "        residual_1 = x\n",
    "        #encoder1\n",
    "        out = self.Econv1(x)\n",
    "        out = self.Ebn1(out)\n",
    "        out = Activation(activation='relu')(out)\n",
    "        out = self.Econv2(out)\n",
    "        out = self.Ebn2(out)\n",
    "        out = Activation(activation='relu')(out)\n",
    "        residual_2 = out\n",
    "\n",
    "        #encoder2\n",
    "        out = self.Econv3(out)\n",
    "        out = self.Ebn3(out)\n",
    "        out = Activation(activation='relu')(out)\n",
    "        out = self.Econv4(out)\n",
    "        out = self.Ebn4(out)\n",
    "        out = Activation(activation='relu')(out)\n",
    "        residual_3 = out\n",
    "        \n",
    "        #RDN_CBAM block\n",
    "        # without concate\n",
    "        out = self.Block1(out)\n",
    "        out = self.Block2(out)\n",
    "        out = self.Block3(out)\n",
    "        out = self.Block4(out)\n",
    "        #out = self.Block5(out)\n",
    "        #out = self.Block6(out)\n",
    "\n",
    "        #decoder2\n",
    "        out = add([out,residual_3])\n",
    "        out = self.Dconv1(out)\n",
    "        out = self.Dbn1(out)\n",
    "        out = Activation(activation='relu')(out)\n",
    "        out = self.Dconv2(out)\n",
    "        out = self.Dbn2(out)\n",
    "        out = Activation(activation='relu')(out)\n",
    "\n",
    "        #decoder1\n",
    "        out = add([out,residual_2])\n",
    "        out = self.Dconv3(out)\n",
    "        out = self.Dbn3(out)\n",
    "        out = Activation(activation='relu')(out)\n",
    "        out = self.Dconv4(out)\n",
    "        out = self.Dbn4(out)\n",
    "        out = Activation(activation='relu')(out)\n",
    "\n",
    "        #final\n",
    "        out = add([out,residual_1])\n",
    "        out = self.conv2(out)\n",
    "        out = self.bn1(out)\n",
    "        out = Activation(activation='relu')(out)\n",
    "        #shape: (None,H,W,T,1)\n",
    "        out = self.ta(out)\n",
    "        out = Lambda(lambda x: K.mean(x, axis = [3], keepdims=False))(out) \n",
    "        out = self.conv3(out)\n",
    "        out = Activation(activation='sigmoid')(out)\n",
    "        #shape: (None,H,W,1)\n",
    "        \n",
    "        return out\n",
    "    \n",
    "    def build(self):\n",
    "         # Initialize the graph\n",
    "        self._is_graph_network = True\n",
    "        self._init_graph_network(\n",
    "            inputs=self.input_layer,\n",
    "            outputs=self.out)\n",
    "    "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "StarN2 = STARN(input_shape=(55, 125, 9, 14)) #shape （H,W,T,feature)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Model: \"starn_1\"\n",
      "__________________________________________________________________________________________________\n",
      " Layer (type)                   Output Shape         Param #     Connected to                     \n",
      "==================================================================================================\n",
      " input_1 (InputLayer)           [(None, 55, 125, 9,  0           []                               \n",
      "                                 14)]                                                             \n",
      "                                                                                                  \n",
      " STARN_Conv1 (Conv3D)           (None, 55, 125, 9,   24256       ['input_1[0][0]']                \n",
      "                                64)                                                               \n",
      "                                                                                                  \n",
      " activation (Activation)        (None, 55, 125, 9,   0           ['STARN_Conv1[0][0]']            \n",
      "                                64)                                                               \n",
      "                                                                                                  \n",
      " STARN_Econv1 (Conv3D)          (None, 55, 125, 9,   110656      ['activation[0][0]']             \n",
      "                                64)                                                               \n",
      "                                                                                                  \n",
      " STARN_Ebn1 (BatchNormalization  (None, 55, 125, 9,   256        ['STARN_Econv1[0][0]']           \n",
      " )                              64)                                                               \n",
      "                                                                                                  \n",
      " activation_1 (Activation)      (None, 55, 125, 9,   0           ['STARN_Ebn1[0][0]']             \n",
      "                                64)                                                               \n",
      "                                                                                                  \n",
      " STARN_Econv2 (Conv3D)          (None, 55, 125, 9,   110656      ['activation_1[0][0]']           \n",
      "                                64)                                                               \n",
      "                                                                                                  \n",
      " STARN_Ebn2 (BatchNormalization  (None, 55, 125, 9,   256        ['STARN_Econv2[0][0]']           \n",
      " )                              64)                                                               \n",
      "                                                                                                  \n",
      " activation_2 (Activation)      (None, 55, 125, 9,   0           ['STARN_Ebn2[0][0]']             \n",
      "                                64)                                                               \n",
      "                                                                                                  \n",
      " STARN_Econv3 (Conv3D)          (None, 55, 125, 9,   110656      ['activation_2[0][0]']           \n",
      "                                64)                                                               \n",
      "                                                                                                  \n",
      " STARN_Ebn3 (BatchNormalization  (None, 55, 125, 9,   256        ['STARN_Econv3[0][0]']           \n",
      " )                              64)                                                               \n",
      "                                                                                                  \n",
      " activation_3 (Activation)      (None, 55, 125, 9,   0           ['STARN_Ebn3[0][0]']             \n",
      "                                64)                                                               \n",
      "                                                                                                  \n",
      " STARN_Econv4 (Conv3D)          (None, 55, 125, 9,   110656      ['activation_3[0][0]']           \n",
      "                                64)                                                               \n",
      "                                                                                                  \n",
      " STARN_Ebn4 (BatchNormalization  (None, 55, 125, 9,   256        ['STARN_Econv4[0][0]']           \n",
      " )                              64)                                                               \n",
      "                                                                                                  \n",
      " activation_4 (Activation)      (None, 55, 125, 9,   0           ['STARN_Ebn4[0][0]']             \n",
      "                                64)                                                               \n",
      "                                                                                                  \n",
      " rdab (RDAB)                    (None, 55, 125, 9,   2352362     ['activation_4[0][0]']           \n",
      "                                64)                                                               \n",
      "                                                                                                  \n",
      " rdab_1 (RDAB)                  (None, 55, 125, 9,   2352362     ['rdab[0][0]']                   \n",
      "                                64)                                                               \n",
      "                                                                                                  \n",
      " rdab_2 (RDAB)                  (None, 55, 125, 9,   2352362     ['rdab_1[0][0]']                 \n",
      "                                64)                                                               \n",
      "                                                                                                  \n",
      " rdab_3 (RDAB)                  (None, 55, 125, 9,   2352362     ['rdab_2[0][0]']                 \n",
      "                                64)                                                               \n",
      "                                                                                                  \n",
      " add (Add)                      (None, 55, 125, 9,   0           ['rdab_3[0][0]',                 \n",
      "                                64)                               'activation_4[0][0]']           \n",
      "                                                                                                  \n",
      " STARN_Dconv1 (Conv3D)          (None, 55, 125, 9,   110656      ['add[0][0]']                    \n",
      "                                64)                                                               \n",
      "                                                                                                  \n",
      " STARN_Dbn1 (BatchNormalization  (None, 55, 125, 9,   256        ['STARN_Dconv1[0][0]']           \n",
      " )                              64)                                                               \n",
      "                                                                                                  \n",
      " activation_5 (Activation)      (None, 55, 125, 9,   0           ['STARN_Dbn1[0][0]']             \n",
      "                                64)                                                               \n",
      "                                                                                                  \n",
      " STARN_Dconv2 (Conv3D)          (None, 55, 125, 9,   110656      ['activation_5[0][0]']           \n",
      "                                64)                                                               \n",
      "                                                                                                  \n",
      " STARN_Dbn2 (BatchNormalization  (None, 55, 125, 9,   256        ['STARN_Dconv2[0][0]']           \n",
      " )                              64)                                                               \n",
      "                                                                                                  \n",
      " activation_6 (Activation)      (None, 55, 125, 9,   0           ['STARN_Dbn2[0][0]']             \n",
      "                                64)                                                               \n",
      "                                                                                                  \n",
      " add_1 (Add)                    (None, 55, 125, 9,   0           ['activation_6[0][0]',           \n",
      "                                64)                               'activation_2[0][0]']           \n",
      "                                                                                                  \n",
      " STARN_Dconv3 (Conv3D)          (None, 55, 125, 9,   110656      ['add_1[0][0]']                  \n",
      "                                64)                                                               \n",
      "                                                                                                  \n",
      " STARN_Dbn3 (BatchNormalization  (None, 55, 125, 9,   256        ['STARN_Dconv3[0][0]']           \n",
      " )                              64)                                                               \n",
      "                                                                                                  \n",
      " activation_7 (Activation)      (None, 55, 125, 9,   0           ['STARN_Dbn3[0][0]']             \n",
      "                                64)                                                               \n",
      "                                                                                                  \n",
      " STARN_Dconv4 (Conv3D)          (None, 55, 125, 9,   110656      ['activation_7[0][0]']           \n",
      "                                64)                                                               \n",
      "                                                                                                  \n",
      " STARN_Dbn4 (BatchNormalization  (None, 55, 125, 9,   256        ['STARN_Dconv4[0][0]']           \n",
      " )                              64)                                                               \n",
      "                                                                                                  \n",
      " activation_8 (Activation)      (None, 55, 125, 9,   0           ['STARN_Dbn4[0][0]']             \n",
      "                                64)                                                               \n",
      "                                                                                                  \n",
      " add_2 (Add)                    (None, 55, 125, 9,   0           ['activation_8[0][0]',           \n",
      "                                64)                               'activation[0][0]']             \n",
      "                                                                                                  \n",
      " STARN_Conv2 (Conv3D)           (None, 55, 125, 9,   1729        ['add_2[0][0]']                  \n",
      "                                1)                                                                \n",
      "                                                                                                  \n",
      " STARN_Bn1 (BatchNormalization)  (None, 55, 125, 9,   4          ['STARN_Conv2[0][0]']            \n",
      "                                1)                                                                \n",
      "                                                                                                  \n",
      " activation_9 (Activation)      (None, 55, 125, 9,   0           ['STARN_Bn1[0][0]']              \n",
      "                                1)                                                                \n",
      "                                                                                                  \n",
      " temporal_attention_4 (Temporal  (None, 55, 125, 9,   180        ['activation_9[0][0]']           \n",
      " Attention)                     1)                                                                \n",
      "                                                                                                  \n",
      " lambda (Lambda)                (None, 55, 125, 1)   0           ['temporal_attention_4[0][0]']   \n",
      "                                                                                                  \n",
      " STARN_Conv3 (Conv2D)           (None, 55, 125, 1)   5           ['lambda[0][0]']                 \n",
      "                                                                                                  \n",
      " activation_10 (Activation)     (None, 55, 125, 1)   0           ['STARN_Conv3[0][0]']            \n",
      "                                                                                                  \n",
      "==================================================================================================\n",
      "Total params: 10,322,918\n",
      "Trainable params: 10,321,892\n",
      "Non-trainable params: 1,026\n",
      "__________________________________________________________________________________________________\n"
     ]
    }
   ],
   "source": [
    "StarN2.summary()"
   ]
  }
 ],
 "metadata": {
  "interpreter": {
   "hash": "12f4ae37f749e26047b6c90ec13e3c0736099234c7c5f03a9452dabd216507b1"
  },
  "kernelspec": {
   "display_name": "Python 3.8.8 ('sgyl_py38')",
   "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.8"
  },
  "orig_nbformat": 4
 },
 "nbformat": 4,
 "nbformat_minor": 2
}