Basic overfit tests now pass meaningfully

  def test_attn_lstm_multitask_classification_overfit(self):

    """Test support graph-conv multitask overfits tiny data."""

    n_tasks = 1

    n_samples = 10

    n_features = 3

    n_test = 5

    n_support = 9

    n_classes = 2

    n_feat = 71

    max_depth = 4

    n_pos = 6

    n_neg = 4

    test_batch_size = 10

    support_batch_size = n_pos + n_neg

    replace = False

    # Load mini log-solubility dataset.

    splittype = "scaffold"

    verbosity = "high"

    classification_metric = Metric(metrics.accuracy_score, verbosity=verbosity)

    n_feat = 71

    batch_size = 10

    test_batch_size = 10

    support_batch_size = 10

    support_model = SequentialSupportGraphModel(n_feat)

    # Add layers

    support_model.join(AttnLSTMEmbedding(test_batch_size, support_batch_size,

                                         max_depth))

    ## Gather Projection

    #support_model.add(Dense(128, activation='relu'))

    #support_model.add_test(BatchNormalization(epsilon=1e-5, mode=1))

    #support_model.add_support(BatchNormalization(epsilon=1e-5, mode=1))

    #support_model.add(GraphGather(batch_size, activation="tanh"))

    with self.test_session() as sess:

      model = SupportGraphClassifier(

        sess, support_model, n_tasks, self.model_dir, 

      # Fit trained model. Dataset has 6 positives and 4 negatives, so set

      # n_pos/n_neg accordingly.  Set replace to false to ensure full dataset

      # is always passed in to support.

      model.fit(dataset, nb_epoch=0, n_trials_per_epoch=10, n_pos=6, n_neg=4,

      # TODO(rbharath): Why does this work with 0 epochs?!!!

      # I think it's because the distance calculation is still meaningful even with

      # random features. The cutoffs also mean that the outputted scores vectors threshold

      # at logit(epsilon), logit(1-epsilon) and stay fixed.

      model.fit(dataset, nb_epoch=10, n_trials_per_epoch=10, n_pos=n_pos, n_neg=n_neg,

                replace=False)

      model.save()

      # model has mastered memorization of provided support.

      scores = model.evaluate(dataset, range(n_tasks),

                              classification_metric, n_trials=5,

                              n_pos=6, n_neg=4,

                              n_pos=n_pos, n_neg=n_neg,

                              exclude_support=False, replace=False)

      print("scores")

      print(scores)

    Computes prediction yhat (eqn (1) in Matching networks) of class for test

    compounds.

    """

    # Get featurization for x

    # Get featurization for test 

    # Shape (n_test, n_feat)

    test_feat = self.model.get_test_output()  

    # Get featurization for support

    # Shape (n_support, n_feat)

    support_feat = self.model.get_support_output()  

    # Computes the inner part c() of the kernel

      support_feat = tf.expand_dims(support_feat, 0)

      max_dist_sq = 20

      g = -tf.maximum(tf.reduce_sum(tf.square(test_feat - support_feat), 2), max_dist_sq)

    # Note that gram matrix g has shape (n_test, n_support)

    # soft corresponds to a(xhat, x_i) in eqn (1) of Matching Networks paper 

    # https://arxiv.org/pdf/1606.04080v1.pdf

    # Computes softmax across axis 1, (so sums distances to support set for

    # each test entry)

    # Shape (n_test, n_support)

    soft = tf.nn.softmax(g)  # Renormalize

    # Weighted sum of support labels

    # Shape (n_support, 1)

    support_labels = tf.expand_dims(self.support_label_placeholder, 1)

    # pred is yhat in eqn (1) of Matching Networks.

    # Shape squeeze((n_test, n_support) * (n_support, 1)) = (n_test,)

    pred = tf.squeeze(tf.matmul(soft, support_labels), [1])

    # Clip softmax probabilities to range [epsilon, 1-epsilon]

    # Shape (n_test,)

    pred = tf.clip_by_value(pred, K.epsilon(), 1.-K.epsilon())

    # Convert to logit space using inverse sigmoid (logit) function

    # logit function: log(pred) - log(1-pred)

    # Used to invoke tf.nn.sigmoid_cross_entropy_with_logits

    # in Cross Entropy calculation.

    # Shape (n_test,)

    scores = tf.log(pred) - tf.log(tf.constant(1., dtype=tf.float32)-pred)

    return pred, scores