Python tensorflow.keras.layers.Concatenate() Examples

def load(input_shape, output_shape, cfg):
    nb_lstm_states = int(cfg['nb_lstm_states'])


    inputs = KL.Input(shape=input_shape)
    x = KL.CuDNNLSTM(units=nb_lstm_states, unit_forget_bias=True)(inputs)

    x = KL.Dense(512)(x)
    x = KL.Activation('relu')(x)
    x = KL.Dropout(0.2)(x)

    x = KL.Dense(256)(x)
    x = KL.Activation('relu')(x)
    x = KL.Dropout(0.3)(x)

    mu = KL.Dense(1)(x)
    std = KL.Dense(1)(x)
    activation_fn = get_activation_function_by_name(cfg['activation_function'])
    std = KL.Activation(activation_fn, name="exponential_activation")(std)

    output = KL.Concatenate(axis=-1)([std, mu])
    model = KM.Model(inputs=[inputs], outputs=[output])

    return model 

def test_single_ddpg_input():
    nb_actions = 2

    actor = Sequential()
    actor.add(Flatten(input_shape=(2, 3)))
    actor.add(Dense(nb_actions))

    action_input = Input(shape=(nb_actions,), name='action_input')
    observation_input = Input(shape=(2, 3), name='observation_input')
    x = Concatenate()([action_input, Flatten()(observation_input)])
    x = Dense(1)(x)
    critic = Model(inputs=[action_input, observation_input], outputs=x)

    memory = SequentialMemory(limit=10, window_length=2)
    agent = DDPGAgent(actor=actor, critic=critic, critic_action_input=action_input, memory=memory,
                      nb_actions=2, nb_steps_warmup_critic=5, nb_steps_warmup_actor=5, batch_size=4)
    agent.compile('sgd')
    agent.fit(MultiInputTestEnv((3,)), nb_steps=10) 

def incorporate_embeddings(self, x):
        """Puts relevant data through embedding layers and then concatenates the result with the rest of the data ready
        to then be put through the hidden layers"""
        all_embedded_data = []
        for embedding_layer_ix, embedding_var in enumerate(self.columns_of_data_to_be_embedded):
            data = x[:, :, embedding_var]
            embedded_data = self.embedding_layers[embedding_layer_ix](data)
            all_embedded_data.append(embedded_data)
        if len(all_embedded_data) > 1: all_embedded_data = Concatenate(axis=2)(all_embedded_data)
        else: all_embedded_data = all_embedded_data[0]
        non_embedded_columns = [col for col in range(x.shape[2]) if col not in self.columns_of_data_to_be_embedded]
        if len(non_embedded_columns) > 0:
            x = tf.gather(x, non_embedded_columns, axis=2)
            x = Concatenate(axis=2)([tf.dtypes.cast(x, float), all_embedded_data])
        else: x = all_embedded_data
        return x 

def process_output_layers(self, x, restricted_to_final_seq):
        """Puts the data x through all the output layers"""
        out = None
        for output_layer_ix, output_layer in enumerate(self.output_layers):
            if type(output_layer) == Dense:
                if self.return_final_seq_only and not restricted_to_final_seq:
                    x = x[:, -1, :]
                    restricted_to_final_seq = True
                temp_output = output_layer(x)
            else:
                temp_output = output_layer(x)
                activation = self.get_activation(self.output_activation, output_layer_ix)
                temp_output = activation(temp_output)
            if out is None: out = temp_output
            else:
                if restricted_to_final_seq: dim = 1
                else: dim = 2
                out = Concatenate(axis=dim)([out, temp_output])
        return out 

def incorporate_embeddings(self, x):
        """Puts relevant data through embedding layers and then concatenates the result with the rest of the data ready
        to then be put through the hidden layers"""
        all_embedded_data = []
        for embedding_layer_ix, embedding_var in enumerate(self.columns_of_data_to_be_embedded):
            data = x[:, embedding_var]
            embedded_data = self.embedding_layers[embedding_layer_ix](data)
            all_embedded_data.append(embedded_data)
        if len(all_embedded_data) > 1: all_embedded_data = Concatenate(axis=1)(all_embedded_data)
        else: all_embedded_data = all_embedded_data[0]
        non_embedded_columns = [col for col in range(x.shape[1]) if col not in self.columns_of_data_to_be_embedded]
        if len(non_embedded_columns) > 0:
            x = tf.gather(x, non_embedded_columns, axis=1)
            x = Concatenate(axis=1)([tf.dtypes.cast(x, float), all_embedded_data])
        else: x = all_embedded_data
        return x 

def build(self, hp, inputs=None):
        inputs = nest.flatten(inputs)
        if len(inputs) == 1:
            return inputs

        merge_type = self.merge_type or hp.Choice('merge_type',
                                                  ['add', 'concatenate'],
                                                  default='add')

        if not all([shape_compatible(input_node.shape, inputs[0].shape) for
                    input_node in inputs]):
            new_inputs = []
            for input_node in inputs:
                new_inputs.append(Flatten().build(hp, input_node))
            inputs = new_inputs

        # TODO: Even inputs have different shape[-1], they can still be Add(
        #  ) after another layer. Check if the inputs are all of the same
        #  shape
        if all([input_node.shape == inputs[0].shape for input_node in inputs]):
            if merge_type == 'add':
                return layers.Add(inputs)

        return layers.Concatenate()(inputs) 

def _build_q_NN(self):
        input_states = Input(shape=(self.observation_size,))
        input_action = Input(shape=(self.action_size,))
        input_layer = Concatenate()([input_states, input_action])
        
        lay1 = Dense(self.observation_size)(input_layer)
        lay1 = Activation('relu')(lay1)
        
        lay2 = Dense(self.observation_size)(lay1)
        lay2 = Activation('relu')(lay2)
        
        lay3 = Dense(2*self.action_size)(lay2)
        lay3 = Activation('relu')(lay3)
        
        advantage = Dense(1, activation = 'linear')(lay3)
        
        model = Model(inputs=[input_states, input_action], outputs=[advantage])
        model.compile(loss='mse', optimizer=Adam(lr=self.lr_))
        
        return model 

def test_ddpg():
    # TODO: replace this with a simpler environment where we can actually test if it finds a solution
    env = gym.make('Pendulum-v0')
    np.random.seed(123)
    env.seed(123)
    random.seed(123)
    nb_actions = env.action_space.shape[0]

    actor = Sequential()
    actor.add(Flatten(input_shape=(1,) + env.observation_space.shape))
    actor.add(Dense(16))
    actor.add(Activation('relu'))
    actor.add(Dense(nb_actions))
    actor.add(Activation('linear'))

    action_input = Input(shape=(nb_actions,), name='action_input')
    observation_input = Input(shape=(1,) + env.observation_space.shape, name='observation_input')
    flattened_observation = Flatten()(observation_input)
    x = Concatenate()([action_input, flattened_observation])
    x = Dense(16)(x)
    x = Activation('relu')(x)
    x = Dense(1)(x)
    x = Activation('linear')(x)
    critic = Model(inputs=[action_input, observation_input], outputs=x)
    
    memory = SequentialMemory(limit=1000, window_length=1)
    random_process = OrnsteinUhlenbeckProcess(theta=.15, mu=0., sigma=.3)
    agent = DDPGAgent(nb_actions=nb_actions, actor=actor, critic=critic, critic_action_input=action_input,
                      memory=memory, nb_steps_warmup_critic=50, nb_steps_warmup_actor=50,
                      random_process=random_process, gamma=.99, target_model_update=1e-3)
    agent.compile([Adam(lr=1e-3), Adam(lr=1e-3)])

    agent.fit(env, nb_steps=400, visualize=False, verbose=0, nb_max_episode_steps=100)
    h = agent.test(env, nb_episodes=2, visualize=False, nb_max_episode_steps=100)
    # TODO: evaluate history 

def yolo2_predictions(feature_maps, feature_channel_nums, num_anchors, num_classes):
    f1, f2 = feature_maps
    f1_channel_num, f2_channel_num = feature_channel_nums

    x1 = compose(
        DarknetConv2D_BN_Leaky(f1_channel_num, (3, 3)),
        DarknetConv2D_BN_Leaky(f1_channel_num, (3, 3)))(f1)

    # Here change the f2 channel number to f2_channel_num//8 first,
    # then expand back to f2_channel_num//2 with "space_to_depth_x2"
    x2 = DarknetConv2D_BN_Leaky(f2_channel_num//8, (1, 1))(f2)
    # TODO: Allow Keras Lambda to use func arguments for output_shape?
    x2_reshaped = Lambda(
        space_to_depth_x2,
        output_shape=space_to_depth_x2_output_shape,
        name='space_to_depth')(x2)

    x = Concatenate()([x2_reshaped, x1])
    x = DarknetConv2D_BN_Leaky(f1_channel_num, (3, 3))(x)
    y = DarknetConv2D(num_anchors * (num_classes + 5), (1, 1), name='predict_conv')(x)

    return y 

def call(self, inputs):
        conv_7x7 = self.conv_7x7(inputs)
        pooled_inputs = self.pool_input(inputs)
        outputs = [pooled_inputs, conv_7x7]
        residual_outputs = []
        for idx in range(self.n_downsample - 1):
            outputs = self.dense_conv[idx](outputs)
            concat_outputs = Concatenate()(outputs)
            outputs = [concat_outputs]

            # Pool each dense layer to match output size
            pooled_outputs = self.pooled_outputs[idx](outputs)
            residual_outputs.append(Concatenate()(pooled_outputs))

            outputs = self.transition_down[idx](outputs)

        outputs = self.dense_conv[-1](outputs)
        outputs = Concatenate()(outputs)
        residual_outputs.append(outputs)
        residual_outputs = [
            Compression(self.compression_factor)(res) for res in residual_outputs
        ]
        outputs = Concatenate()(residual_outputs)
        return [outputs] 

def resblock_body(x, num_filters, num_blocks, all_narrow=True):
    '''A series of resblocks starting with a downsampling Convolution2D'''
    # Darknet uses left and top padding instead of 'same' mode
    x = ZeroPadding2D(((1,0),(1,0)))(x)
    x = DarknetConv2D_BN_Mish(num_filters, (3,3), strides=(2,2))(x)

    res_connection = DarknetConv2D_BN_Mish(num_filters//2 if all_narrow else num_filters, (1,1))(x)
    x = DarknetConv2D_BN_Mish(num_filters//2 if all_narrow else num_filters, (1,1))(x)

    for i in range(num_blocks):
        y = compose(
                DarknetConv2D_BN_Mish(num_filters//2, (1,1)),
                DarknetConv2D_BN_Mish(num_filters//2 if all_narrow else num_filters, (3,3)))(x)
        x = Add()([x,y])

    x = DarknetConv2D_BN_Mish(num_filters//2 if all_narrow else num_filters, (1,1))(x)
    x = Concatenate()([x , res_connection])

    return DarknetConv2D_BN_Mish(num_filters, (1,1))(x) 

def residual_block_id(self,tensor, feature_n,name=None):
        if name != None:
            depconv_1  = DepthwiseConv2D(3,2,padding='same',name=name+"/dconv")(tensor)
            conv_2     = Conv2D(feature_n,1,name=name+"/conv")(depconv_1)
        else:
            depconv_1  = DepthwiseConv2D(3,2,padding='same')(tensor)
            conv_2     = Conv2D(feature_n,1)(depconv_1)


        maxpool_1  = MaxPool2D(pool_size=(2,2),strides=(2,2),padding='same')(tensor)
        conv_zeros = Conv2D(feature_n/2,2,strides=2,use_bias=False,kernel_initializer=tf.zeros_initializer())(tensor)

        padding_1  = Concatenate(axis=-1)([maxpool_1,conv_zeros])#self.feature_padding(maxpool_1)

        add = Add()([padding_1,conv_2])
        relu = ReLU()(add)

        return relu
    
    #def feature_padding(self,tensor,channels_n=0):
    #    #pad = tf.keras.layers.ZeroPadding2D(((0,0),(0,0),(0,tensor.shape[3])))(tensor)
    #    return Concatenate(axis=3)([tensor,pad]) 

def __init__(self, growth_rate=64, bottleneck_factor=1, **kwargs):
        # super(DenseConv2D, self).__init__(self, **kwargs)
        self.concat = Concatenate()

        bottleneck_filters = int(np.round(growth_rate * bottleneck_factor))

        self.bottleneck_1x1 = layers.Conv2D(
            bottleneck_filters,
            (1, 1),
            padding="same",
            activation="selu",
            kernel_initializer="lecun_normal",
        )
        self.conv_3x3 = layers.Conv2D(
            growth_rate,
            (3, 3),
            padding="same",
            activation="selu",
            kernel_initializer="lecun_normal",
        ) 

def call(self, x):
        x0_0 = self.conv0_0(x)
        x1_0 = self.conv1_0(self.pool(x0_0))
        x0_1 = self.conv0_1(Concatenate()([x0_0, self.Up(x1_0)]))

        x2_0 = self.conv2_0(self.pool(x1_0))
        x1_1 = self.conv1_1(Concatenate()([x1_0, self.Up(x2_0)]))
        x0_2 = self.conv0_2(Concatenate()([x0_0, x0_1, self.Up(x1_1)]))

        x3_0 = self.conv3_0(self.pool(x2_0))
        x2_1 = self.conv2_1(Concatenate()([x2_0, self.Up(x3_0)]))
        x1_2 = self.conv1_2(Concatenate()([x1_0, x1_1, self.Up(x2_1)]))
        x0_3 = self.conv0_3(Concatenate()([x0_0, x0_1, x0_2, self.Up(x1_2)]))

        x4_0 = self.conv4_0(self.pool(x3_0))
        x3_1 = self.conv3_1(Concatenate()([x3_0, self.Up(x4_0)]))
        x2_2 = self.conv2_2(Concatenate()([x2_0, x2_1, self.Up(x3_1)]))
        x1_3 = self.conv1_3(Concatenate()([x1_0, x1_1, x1_2, self.Up(x2_2)]))
        x0_4 = self.conv0_4(Concatenate()([x0_0, x0_1, x0_2, x0_3, self.Up(x1_3)]))

        output = self.final(x0_4)
        return output 

def __init__(self, n_output_channels, activation="linear", name=None, **kwargs):
        self.output_channels = layers.Conv2D(
            n_output_channels, (1, 1), padding="same", activation=activation, name=name
        )
        self.concat = Concatenate() 

def shuffle_unit(inputs, out_channels, bottleneck_ratio,strides=2,stage=1,block=1):
    if K.image_data_format() == 'channels_last':
        bn_axis = -1
    else:
        raise ValueError('Only channels last supported')

    prefix = 'stage{}/block{}'.format(stage, block)
    bottleneck_channels = int(out_channels * bottleneck_ratio)
    if strides < 2:
        c_hat, c = channel_split(inputs, '{}/spl'.format(prefix))
        inputs = c

    x = Conv2D(bottleneck_channels, kernel_size=(1,1), strides=1, padding='same', name='{}/1x1conv_1'.format(prefix))(inputs)
    x = BatchNormalization(axis=bn_axis, name='{}/bn_1x1conv_1'.format(prefix))(x)
    x = Activation('relu', name='{}/relu_1x1conv_1'.format(prefix))(x)
    x = DepthwiseConv2D(kernel_size=3, strides=strides, padding='same', name='{}/3x3dwconv'.format(prefix))(x)
    x = BatchNormalization(axis=bn_axis, name='{}/bn_3x3dwconv'.format(prefix))(x)
    x = Conv2D(bottleneck_channels, kernel_size=1,strides=1,padding='same', name='{}/1x1conv_2'.format(prefix))(x)
    x = BatchNormalization(axis=bn_axis, name='{}/bn_1x1conv_2'.format(prefix))(x)
    x = Activation('relu', name='{}/relu_1x1conv_2'.format(prefix))(x)

    if strides < 2:
        ret = Concatenate(axis=bn_axis, name='{}/concat_1'.format(prefix))([x, c_hat])
    else:
        s2 = DepthwiseConv2D(kernel_size=3, strides=2, padding='same', name='{}/3x3dwconv_2'.format(prefix))(inputs)
        s2 = BatchNormalization(axis=bn_axis, name='{}/bn_3x3dwconv_2'.format(prefix))(s2)
        s2 = Conv2D(bottleneck_channels, kernel_size=1,strides=1,padding='same', name='{}/1x1_conv_3'.format(prefix))(s2)
        s2 = BatchNormalization(axis=bn_axis, name='{}/bn_1x1conv_3'.format(prefix))(s2)
        s2 = Activation('relu', name='{}/relu_1x1conv_3'.format(prefix))(s2)
        ret = Concatenate(axis=bn_axis, name='{}/concat_2'.format(prefix))([x, s2])

    ret = Lambda(channel_shuffle, name='{}/channel_shuffle'.format(prefix))(ret)

    return ret 

def call(self, x):
        e1 = self.RRCNN1(x)

        e2 = self.Maxpool(e1)
        e2 = self.RRCNN2(e2)

        e3 = self.Maxpool1(e2)
        e3 = self.RRCNN3(e3)

        e4 = self.Maxpool2(e3)
        e4 = self.RRCNN4(e4)

        e5 = self.Maxpool3(e4)
        e5 = self.RRCNN5(e5)

        d5 = self.Up5(e5)
        d5 = Concatenate()([e4, d5])
        d5 = self.Up_RRCNN5(d5)

        d4 = self.Up4(d5)
        d4 = Concatenate()([e3, d4])
        d4 = self.Up_RRCNN4(d4)

        d3 = self.Up3(d4)
        d3 = Concatenate()([e2, d3])
        d3 = self.Up_RRCNN3(d3)

        d2 = self.Up2(d3)
        d2 = Concatenate()([e1, d2])
        d2 = self.Up_RRCNN2(d2)

        out = self.Conv(d2)
        return out 

def yolo3_predictions(feature_maps, feature_channel_nums, num_anchors, num_classes, use_spp=False):
    f1, f2, f3 = feature_maps
    f1_channel_num, f2_channel_num, f3_channel_num = feature_channel_nums

    #feature map 1 head & output (13x13 for 416 input)
    if use_spp:
        x, y1 = make_spp_last_layers(f1, f1_channel_num//2, num_anchors * (num_classes + 5), predict_id='1')
    else:
        x, y1 = make_last_layers(f1, f1_channel_num//2, num_anchors * (num_classes + 5), predict_id='1')

    #upsample fpn merge for feature map 1 & 2
    x = compose(
            DarknetConv2D_BN_Leaky(f2_channel_num//2, (1,1)),
            UpSampling2D(2))(x)
    x = Concatenate()([x,f2])

    #feature map 2 head & output (26x26 for 416 input)
    x, y2 = make_last_layers(x, f2_channel_num//2, num_anchors*(num_classes+5), predict_id='2')

    #upsample fpn merge for feature map 2 & 3
    x = compose(
            DarknetConv2D_BN_Leaky(f3_channel_num//2, (1,1)),
            UpSampling2D(2))(x)
    x = Concatenate()([x, f3])

    #feature map 3 head & output (52x52 for 416 input)
    x, y3 = make_last_layers(x, f3_channel_num//2, num_anchors*(num_classes+5), predict_id='3')

    return y1, y2, y3 

def dense_block(x, nb_layers, nb_channels, growth_rate, dropout_rate=None, bottleneck=False, weight_decay=1e-4):
    """
    Creates a dense block and concatenates inputs
    """
    
    x_list = [x]
    for i in range(nb_layers):
        cb = convolution_block(x, growth_rate, dropout_rate, bottleneck, weight_decay)
        x_list.append(cb)
        x = Concatenate(axis=-1)(x_list)
        nb_channels += growth_rate
    return x, nb_channels 

def call(self, x):
        e1 = self.Conv1(x)

        e2 = self.Maxpool1(e1)
        e2 = self.Conv2(e2)

        e3 = self.Maxpool2(e2)
        e3 = self.Conv3(e3)

        e4 = self.Maxpool3(e3)
        e4 = self.Conv4(e4)

        e5 = self.Maxpool4(e4)
        e5 = self.Conv5(e5)

        d5 = self.Up5(e5)
        d5 = Concatenate()([e4, d5])

        d5 = self.Up_conv5(d5)

        d4 = self.Up4(d5)
        d4 = Concatenate()([e3, d4])
        d4 = self.Up_conv4(d4)

        d3 = self.Up3(d4)
        d3 = Concatenate()([e2, d3])
        d3 = self.Up_conv3(d3)

        d2 = self.Up2(d3)
        d2 = Concatenate()([e1, d2])
        d2 = self.Up_conv2(d2)

        out = self.Conv(d2)
        return out 

def tiny_yolo4lite_predictions(feature_maps, feature_channel_nums, num_anchors, num_classes, use_spp):
    f1, f2 = feature_maps
    f1_channel_num, f2_channel_num = feature_channel_nums

    #feature map 1 head (13 x 13 x f1_channel_num//2 for 416 input)
    x1 = DarknetConv2D_BN_Leaky(f1_channel_num//2, (1,1))(f1)
    if use_spp:
        x1 = Spp_Conv2D_BN_Leaky(x1, f1_channel_num//2)

    #upsample fpn merge for feature map 1 & 2
    x1_upsample = compose(
            DarknetConv2D_BN_Leaky(f2_channel_num//2, (1,1)),
            UpSampling2D(2))(x1)
    x2 = compose(
            Concatenate(),
            #DarknetConv2D_BN_Leaky(f2_channel_num, (3,3)),
            Depthwise_Separable_Conv2D_BN_Leaky(filters=f2_channel_num, kernel_size=(3, 3), block_id_str='pred_1'))([x1_upsample, f2])

    #feature map 2 output (26 x 26 x f2_channel_num for 416 input)
    y2 = DarknetConv2D(num_anchors*(num_classes+5), (1,1), name='predict_conv_2')(x2)

    #downsample fpn merge for feature map 2 & 1
    x2_downsample = compose(
            ZeroPadding2D(((1,0),(1,0))),
            #DarknetConv2D_BN_Leaky(f1_channel_num//2, (3,3), strides=(2,2)),
            Darknet_Depthwise_Separable_Conv2D_BN_Leaky(f1_channel_num//2, (3,3), strides=(2,2), block_id_str='pred_2'))(x2)
    x1 = compose(
            Concatenate(),
            #DarknetConv2D_BN_Leaky(f1_channel_num, (3,3)),
            Depthwise_Separable_Conv2D_BN_Leaky(filters=f1_channel_num, kernel_size=(3, 3), block_id_str='pred_3'))([x2_downsample, x1])

    #feature map 1 output (13 x 13 x f1_channel_num for 416 input)
    y1 = DarknetConv2D(num_anchors*(num_classes+5), (1,1), name='predict_conv_1')(x1)

    return y1, y2 

def call(self, x):
        e1 = self.Conv1(x)

        e2 = self.Maxpool1(e1)
        e2 = self.Conv2(e2)

        e3 = self.Maxpool2(e2)
        e3 = self.Conv3(e3)

        e4 = self.Maxpool3(e3)
        e4 = self.Conv4(e4)

        e5 = self.Maxpool4(e4)
        e5 = self.Conv5(e5)

        d5 = self.Up5(e5)
        x4 = self.Att5([d5,e4])
        d5 = Concatenate()([x4, d5])
        d5 = self.Up_conv5(d5)

        d4 = self.Up4(d5)
        x3 = self.Att4([d4,e3])
        d4 = Concatenate()([x3, d4])
        d4 = self.Up_conv4(d4)

        d3 = self.Up3(d4)
        x2 = self.Att3([d3,e2])
        d3 = Concatenate()([x2, d3])
        d3 = self.Up_conv3(d3)

        d2 = self.Up2(d3)
        x1 = self.Att2([d2,e1])
        d2 = Concatenate()([x1, d2])
        d2 = self.Up_conv2(d2)

        out = self.Conv(d2)
        return out 

def Spp_Conv2D_BN_Leaky(x, num_filters):
    y1 = MaxPooling2D(pool_size=(5,5), strides=(1,1), padding='same')(x)
    y2 = MaxPooling2D(pool_size=(9,9), strides=(1,1), padding='same')(x)
    y3 = MaxPooling2D(pool_size=(13,13), strides=(1,1), padding='same')(x)

    y = compose(
            Concatenate(),
            DarknetConv2D_BN_Leaky(num_filters, (1,1)))([y1, y2, y3, x])
    return y 

def additive_self_attention(units, n_hidden=None, n_output_features=None, activation=None):
    """
    Compute additive self attention for time series of vectors (with batch dimension)
            the formula: score(h_i, h_j) = <v, tanh(W_1 h_i + W_2 h_j)>
            v is a learnable vector of n_hidden dimensionality,
            W_1 and W_2 are learnable [n_hidden, n_input_features] matrices

    Args:
        units: tf tensor with dimensionality [batch_size, time_steps, n_input_features]
        n_hidden: number of2784131 units in hidden representation of similarity measure
        n_output_features: number of features in output dense layer
        activation: activation at the output

    Returns:
        output: self attended tensor with dimensionality [batch_size, time_steps, n_output_features]
        """
    n_input_features = K.int_shape(units)[2]
    if n_hidden is None:
        n_hidden = n_input_features
    if n_output_features is None:
        n_output_features = n_input_features
    exp1 = Lambda(lambda x: expand_tile(x, axis=1))(units)
    exp2 = Lambda(lambda x: expand_tile(x, axis=2))(units)
    units_pairs = Concatenate(axis=3)([exp1, exp2])
    query = Dense(n_hidden, activation="tanh")(units_pairs)
    attention = Dense(1, activation=lambda x: softmax(x, axis=2))(query)
    attended_units = Lambda(lambda x: K.sum(attention * x, axis=2))(exp1)
    output = Dense(n_output_features, activation=activation)(attended_units)
    return output 

def call(self, inputs):
        residual_outputs = [Concatenate()(inputs)]
        outputs = self.transition_input(inputs)

        # Encoder
        for idx in range(self.n_downsample - 1):
            outputs = self.dense_conv_down[idx](outputs)
            concat_outputs = Concatenate()(outputs)
            outputs = [concat_outputs]
            residual_outputs.append(concat_outputs)
            outputs = self.transition_down[idx](outputs)
        residual_outputs.append(Concatenate()(outputs))
        outputs = self.dense_conv_encoded(outputs)

        # Compress the feature maps for residual connections
        residual_outputs = residual_outputs[::-1]
        residual_outputs = [
            Compression(self.compression_factor)(res) for res in residual_outputs
        ]

        # Decoder
        for idx in range(self.n_upsample):
            outputs.append(residual_outputs[idx])
            outputs = self.dense_conv_up[idx](outputs)
            outputs = self.transition_up[idx](outputs)
        outputs.append(residual_outputs[-1])
        outputs = self.dense_conv_output(outputs)
        return [Concatenate()(outputs)] 

def __init__(self, compression_factor=0.5, **kwargs):
        # super(TransitionDown, self).__init__(self, **kwargs)
        self.concat = Concatenate()
        self.compression_factor = compression_factor

        self.upsample = (
            SubPixelUpscaling()
        )  # layers.UpSampling2D(interpolation='bilinear') 

def __init__(self, compression_factor=0.5, **kwargs):
        # super(Compression, self).__init__(self, **kwargs)
        self.concat = Concatenate()
        self.compression_factor = compression_factor 

def __call__(self, inputs):
        return self.call(inputs)

    # def get_config(self):
    #    config = {}
    #    base_config = super(Concatenate, self).get_config()
    #    return dict(list(base_config.items()) + list(config.items())) 

def __init__(self, **kwargs):
        # super(Concatenate, self).__init__(self, **kwargs)
        self.concat = layers.Concatenate() 

def cvpr2018_net(vol_size, enc_nf, dec_nf, indexing='ij', name="voxelmorph"):
    """
    From https://github.com/voxelmorph/voxelmorph.

    unet architecture for voxelmorph models presented in the CVPR 2018 paper.
    You may need to modify this code (e.g., number of layers) to suit your project needs.

    :param vol_size: volume size. e.g. (256, 256, 256)
    :param enc_nf: list of encoder filters. right now it needs to be 1x4.
           e.g. [16,32,32,32]
    :param dec_nf: list of decoder filters. right now it must be 1x6 (like voxelmorph-1) or 1x7 (voxelmorph-2)
    :return: the keras model
    """
    import tensorflow.keras.layers as KL

    ndims = len(vol_size)
    assert ndims==3, "ndims should be 3. found: %d" % ndims

    src = Input(vol_size + (1,), name='input_src')
    tgt = Input(vol_size + (1,), name='input_tgt')

    input_stack = Concatenate(name='concat_inputs')([src, tgt])

    # get the core model
    x = unet3D(input_stack, img_shape=vol_size, out_im_chans=ndims, nf_enc=enc_nf, nf_dec=dec_nf)

    # transform the results into a flow field.
    Conv = getattr(KL, 'Conv%dD' % ndims)
    flow = Conv(ndims, kernel_size=3, padding='same', name='flow',
                  kernel_initializer=RandomNormal(mean=0.0, stddev=1e-5))(x)

    # warp the source with the flow
    y = SpatialTransformer(interp_method='linear', indexing=indexing)([src, flow])
    # prepare model
    model = Model(inputs=[src, tgt], outputs=[y, flow], name=name)
    return model


##############################################################################
# Appearance transform model
##############################################################################