title: Tensorflow Keras implementation of multi-instance learning problem
body: <p>I'm trying to implement a multi-instance learning model with TensorFlow Keras API. Suppose the input tensor shape is <code>(None, 18, 10, 300)</code>, and I need to perform multi-instance learning with <code>submodel</code> along <code>axis=1</code>. Thus, the input to <code>submodel</code> should be <code>(None, 10, 300)</code>. Suppose the output tensor's shape for a single <code>submodel</code> is <code>(None, 100)</code>, the output of <code>submodel</code>s is then concatenated thus create output shape <code>(None, 18, 100)</code>. </p>
  
  <p>I implemented this mechanism with 2 ways, they should be functionally equal, however, when training the model, the second implementation's accuracy is 20% higher than the first implementation (though both is not good enough). I wonder whether there is functional difference between the two implementation? If both of these implementation is not what I mean to do, could you please give a correct implementation?</p>
  
  <p>Implementation 1. This implementation split the input tensor with a keras <code>Lambda</code> layer, execute the <code>submodel</code>, expand the result's dimension in <code>axis=1</code>, then concatenate the result along <code>axis=1</code>.</p>
  
  <pre class="lang-py prettyprint-override"><code>def get_branch_model(input_shape, submodel, args={}):
      model_input = tf.keras.Input(input_shape)
      sliced_inputs = [tf.keras.layers.Lambda(lambda x: x[:,i])(model_input) 
                       for i in range(input_shape[0])]
      sub_instance = submodel(**args)
      branch_models = [sub_instance(sliced_inputs[i]) for i in range(input_shape[0])]
      expand_layer = tf.keras.layers.Lambda(lambda x: tf.keras.backend.expand_dims(x, axis=1))
      expanded_outputs = [expand_layer(branch_models[i]) for i in range(input_shape[0])]
      concated_layer = tf.keras.layers.Concatenate(axis=1)(expanded_outputs)
      return tf.keras.Model(model_input, concated_layer)
  </code></pre>
  
  <p>Implementation 2. This implementation use <code>tf.slice</code> to split the input tensor, execute the <code>submodel</code>, expand the result's dimension in <code>axis=1</code>, then concatenate the result along <code>axis=1</code>.</p>
  
  <pre class="lang-py prettyprint-override"><code>def __get_filter_layer(total_dim, target_dim, index):
      def tensor_filter(tensor_in):
          nonlocal index
          begin = [0 if i != target_dim else index for i in range(total_dim)]
          size = [-1 if i != target_dim else 1 for i in range(total_dim)]
          return tf.squeeze(tf.slice(tensor_in, begin, size), axis=target_dim)
      return tf.keras.models.Sequential([
          tf.keras.layers.Lambda(tensor_filter)
      ])
  
  def get_branch_model(input_shape, branch_index, output_shape, submodel, args={}):
      model_input = tf.keras.Input(input_shape)
      sliced_inputs = [__get_filter_layer(len(input_shape) + 1, branch_index, i)(model_input) 
                       for i in range(input_shape[branch_index - 1])]
      sub_instance = submodel(**args)
      branch_models = [sub_instance(sliced_inputs[i]) 
                       for i in range(input_shape[branch_index - 1])]
      expand_layer = tf.keras.layers.Lambda(lambda x: tf.keras.backend.expand_dims(x, axis=1))
      expanded_outputs = [expand_layer(branch_models[i]) for i in range(input_shape[0])]
      concated_layer = tf.keras.layers.Concatenate(axis=1)(expanded_outputs)
      return tf.keras.Model(model_input, concated_layer)
  </code></pre>
  
  <blockquote>
    <p>The input for argument <code>input_shape</code> is <code>(18, 10, 300)</code>, and <code>branch_index</code> is <code>1</code>. <code>submodel</code> is a Sequential model, such as <code>tf.keras.Sequential([tf.keras.layers.Dense(...)])</code>.</p>
  </blockquote>