Overfit test runs but fails (save/load issues I believe)

import numpy as np

import unittest

import sklearn

import shutil

import tensorflow as tf

from keras import backend as K

from keras.layers import Dense, BatchNormalization

from tensorflow.python.framework import test_util

from deepchem.featurizers.featurize import DataLoader

from deepchem.featurizers.fingerprints import CircularFingerprint

from sklearn.ensemble import RandomForestClassifier

from deepchem.models.tf_keras_models.keras_layers import GraphPool

from deepchem.models.tf_keras_models.keras_layers import GraphGather

from deepchem.featurizers.graph_features import ConvMolFeaturizer

from multitask_classifier import MultitaskGraphClassifier

from deepchem.models.tf_keras_models.multitask_classifier import MultitaskGraphClassifier

class TestOverfitAPI(TestAPI):

class TestOverfitAPI(test_util.TensorFlowTestCase):

  """

  Test that models can overfit simple datasets.

  """

  def setUp(self):

    super(TestOverfitAPI, self).setUp()

    self.root = '/tmp'

    self.smiles_field = "smiles"

    self.current_dir = os.path.dirname(os.path.abspath(__file__))

    self.train_dir = tempfile.mkdtemp()

    self.data_dir = tempfile.mkdtemp()

    self.model_dir = tempfile.mkdtemp()

  def tearDown(self):

    shutil.rmtree(self.train_dir)

    shutil.rmtree(self.data_dir)

  def test_sklearn_regression_overfit(self):

    """Test that sklearn models can overfit simple regression datasets."""

  def test_graph_conv_multitask_classification_overfit(self):

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

    n_tasks = 10

    n_tasks = 1

    n_samples = 10

    n_features = 3

    n_classes = 2

    n_atoms = 50

    n_feat = 71

    batch_size = 20

    batch_size = 10

    graph_model = SequentialGraphModel(n_atoms, n_feat, batch_size)

    graph_model.add(GraphConv(64, activation='relu'))

    graph_model.add(BatchNormalization(epsilon=1e-5, mode=1))

    graph_model.add(BatchNormalization(epsilon=1e-5, mode=1))

    graph_model.add(GraphGather(batch_size, activation="tanh"))

    model = MultitaskGraphClassifier(

      sess, graph_model, n_tasks, learning_rate=1e-3,

    with self.test_session() as sess:

      tensorflow_model = MultitaskGraphClassifier(

        sess, graph_model, n_tasks, self.model_dir, learning_rate=1e-3,

        learning_rate_decay_time=1000, optimizer_type="adam", beta1=.9,

        beta2=.999, verbosity="high")

      model = TensorflowModel(tensorflow_model, self.model_dir)

      # Fit trained model

      model.fit(dataset)

                          graph_topology=self.graph_topology)

  '''

  def get_batch_size():

  def get_batch_size(self):

    return self.batch_size

  def get_graph_topology():

  def get_graph_topology(self):

    return self.graph_topology

  def get_num_output_features(self):

import os

import sys

import numpy as np

import tensorflow as tf

from keras.layers import Input, Dense

from keras import initializations, activations

from keras import backend as K

from utils import merge_dicts

from deepchem.datasets import pad_features

from deepchem.utils.save import log

from deepchem.models import Model

from deepchem.models.tensorflow_models import TensorflowClassifier

from deepchem.models.tensorflow_models import model_ops

# TODO(rbharath): Find a way to get rid of this import?

from deepchem.models.tf_keras_models.graph_topology import merge_dicts

def get_loss_fn(final_loss):

  # Obtain appropriate loss function

      return tf.reduce_sum(weighted_costs)

  return loss_fn

class MultitaskGraphClassifier(Model):

class MultitaskGraphClassifier(TensorflowClassifier):

  def __init__(self, sess, model, n_tasks,

  def __init__(self, sess, model, n_tasks, logdir,

               final_loss='cross_entropy', learning_rate=.001,

               optimizer_type="adam", learning_rate_decay_time=1000,

               beta1=.9, beta2=.999, verbosity=None):

    self.n_tasks = n_tasks

    self.final_loss = final_loss

    self.model = model 

    self.logdir = logdir

    # Extract model info 

    self.batch_size = self.model.get_batch_size()

    self.init_fn = tf.initialize_all_variables()

    sess.run(self.init_fn)  

    # Path to save checkpoint files, which matches the

    # replicated supervisor's default path.

    self._save_path = os.path.join(logdir, 'model.ckpt')

  def build(self):

    # Create target inputs

    self.label_placeholder = Input(tensor=K.placeholder(

      shape=(None,self.n_tasks), name="label", dtype='bool'))

      shape=(None,self.n_tasks), name="label_placeholder", dtype='bool'))

    ######################################### DEBUG

    print("self.n_tasks")

    print(self.n_tasks)

    ######################################### DEBUG

    self.weight_placeholder = Input(tensor=K.placeholder(

          shape=(None,self.n_tasks), name="weight", dtype='float32'))

          shape=(None,self.n_tasks), name="weight_placholder", dtype='float32'))

    # Create final dense layer from keras 

    feat = self.model.return_outputs()

    # num_dat_dict = {self.num_datapoints_placeholder : self.}

    # Get other optimizer information

    keras_dict = {K.learning_phase() : training}

    # TODO(rbharath): Figure out how to handle phase appropriately

    #keras_dict = {K.learning_phase() : training}

    keras_dict = {}

    feed_dict = merge_dicts([targets_dict, atoms_dict,

                             keras_dict])

    return feed_dict

        softmax.append(tf.nn.softmax(logits, name='softmax_%d' % i))

    return softmax

  def fit(self, dataset, nb_epoch=10, batch_size=50, pad_batches=False,

          log_every_N_batches=50, **kwargs):

          max_checkpoints_to_keep=5, log_every_N_batches=50, **kwargs):

    # Perform the optimization

    log("Training for %d epochs" % nb_epoch, self.verbosity)

    # Save an initial checkpoint.

    saver = tf.train.Saver(max_to_keep=max_checkpoints_to_keep)

    saver.save(self.sess, self._save_path, global_step=0)

    for epoch in range(nb_epoch):

      # TODO(rbharath): This decay shouldn't be hard-coded.

      lr = self.learning_rate / (1 + float(epoch) / self.T)

        self.sess.run(

            self.train_op,

            feed_dict=self.construct_feed_dict(X_b, y_b, w_b))

      saver.save(self.sess, self._save_path, global_step=epoch)

    # Always save a final checkpoint when complete.

    saver.save(self.sess, self._save_path, global_step=epoch+1)

  def save(self):

    """

    No-op since models save themselves during fit()

    """

    pass

  def predict(self, dataset, transformers=[], **kwargs):

    """Wraps predict to set batch_size/padding."""