Updated to MAML code

import deepchem.data

import deepchem.feat

import deepchem.hyper

import deepchem.metalearning

import deepchem.metrics

import deepchem.models

import deepchem.nn

import tensorflow as tf

import time

class MetaLearner(object):

  """Model and data to which the MAML algorithm can be applied.

  data for training it on a large (possibly infinite) set of different tasks.

  """

  def __init__(self):

    self.batch_size = 10

    self.tg = dc.models.TensorGraph(use_queue=False)

    self.features = layers.Feature(shape=(None, 1))

    self.labels = layers.Label(shape=(None, 1))

    hidden1 = layers.Dense(in_layers=self.features, out_channels=40, activation_fn=tf.nn.relu)

    hidden2 = layers.Dense(in_layers=hidden1, out_channels=40, activation_fn=tf.nn.relu)

    output = layers.Dense(in_layers=hidden2, out_channels=1)

    loss = layers.L2Loss(in_layers=[output, self.labels])

    self.tg.add_output(output)

    self.tg.set_loss(loss)

    with self.tg._get_tf("Graph").as_default():

      self.tg.build()

      self._variables = tf.get_collection(tf.GraphKeys.TRAINABLE_VARIABLES)

      self._loss = self._variables[0]+self._variables[1]

  @property

  def loss(self):

    """Get the model's loss function, represented as a Layer or Tensor."""

  and data for your learning problem.  Pass it to a MAML object and call fit().

  """

  def __init__(self, learner, learning_rate=0.001, optimization_steps=1, optimizer=Adam(), model_dir=None):

  def __init__(self,

               learner,

               learning_rate=0.001,

               optimization_steps=1,

               optimizer=Adam(),

               model_dir=None):

    """Create an object for performing meta-optimization.

    Parameters

      updated_variables = learner.variables

      for i in range(optimization_steps):

        gradients = tf.gradients(updated_loss, updated_variables)

        updated_variables = [v if g is None else v-self._learning_rate*g for v,g in zip(updated_variables, gradients)]

        replacements = dict((tf.convert_to_tensor(v1), v2) for v1,v2 in zip(learner.variables, updated_variables))

        updated_variables = [

            v if g is None else v - self._learning_rate * g

            for v, g in zip(updated_variables, gradients)

        ]

        replacements = dict(

            (tf.convert_to_tensor(v1), v2)

            for v1, v2 in zip(learner.variables, updated_variables))

        if i == optimization_steps - 1:

          # In the final loss, use different placeholders for all inputs so the loss will be

          # computed from a different batch.

          for p in self._meta_placeholders:

            replacements[p] = self._meta_placeholders[p]

        updated_loss = tf.contrib.graph_editor.graph_replace(self._loss, replacements)

        updated_loss = tf.contrib.graph_editor.graph_replace(

            self._loss, replacements)

      self._meta_loss = updated_loss

      # Create the optimizer for meta-optimization.

      self._session.run(tf.global_variables_initializer())

      if restore:

        self.restore()

      saver = tf.train.Saver(self.learner.variables, max_to_keep=max_checkpoints_to_keep)

      saver = tf.train.Saver(

          self.learner.variables, max_to_keep=max_checkpoints_to_keep)

      checkpoint_index = 0

      checkpoint_time = time.time()

        feed_dict[self._global_step] = i

        for key, value in self.learner.get_batch().items():

          feed_dict[self._meta_placeholders[key]] = value

        #if i%100 == 0:

        #  print(i, sum(self._session.run(self._loss, feed_dict=feed_dict)), sum(self._session.run(self._meta_loss, feed_dict=feed_dict)))

        self._session.run(self._train_op, feed_dict=feed_dict)

        # Do checkpointing.

        if i == steps-1 or time.time() >= checkpoint_time + checkpoint_interval:

        if i == steps - 1 or time.time(

        ) >= checkpoint_time + checkpoint_interval:

          saver.save(

              self._session,

              self.save_file,

              global_step=checkpoint_index)

              self._session, self.save_file, global_step=checkpoint_index)

          checkpoint_index += 1

          checkpoint_time = time.time()

        self.tg = dc.models.TensorGraph(use_queue=False)

        self.features = Feature(shape=(None, 1))

        self.labels = Label(shape=(None, 1))

        hidden1 = Dense(in_layers=self.features, out_channels=40, activation_fn=tf.nn.relu)

        hidden2 = Dense(in_layers=hidden1, out_channels=40, activation_fn=tf.nn.relu)

        hidden1 = Dense(

            in_layers=self.features, out_channels=40, activation_fn=tf.nn.relu)

        hidden2 = Dense(

            in_layers=hidden1, out_channels=40, activation_fn=tf.nn.relu)

        output = Dense(in_layers=hidden2, out_channels=1)

        loss = L2Loss(in_layers=[output, self.labels])

        self.tg.add_output(output)

        x = np.random.uniform(-5.0, 5.0, (self.batch_size, 1))

        feed_dict = {}

        feed_dict[self.features.out_tensor] = x

        feed_dict[self.labels.out_tensor] = self.amplitude*np.sin(x+self.phase)

        feed_dict[self.labels.out_tensor] = self.amplitude * np.sin(

            x + self.phase)

        return feed_dict

    # Optimize it.

      feed_dict = learner.get_batch()

      for key, value in learner.get_batch().items():

        feed_dict[maml._meta_placeholders[key]] = value

      loss1.append(np.average(np.sqrt(maml._session.run(maml._loss, feed_dict=feed_dict))))

      loss2.append(np.average(np.sqrt(maml._session.run(maml._meta_loss, feed_dict=feed_dict))))

    assert np.average(loss1) > 1.0 # Initially the model should do a bad job of fitting the sine function

    assert np.average(loss2) < 1.0 # After one step of optimization it should do much better

      loss1.append(

          np.average(

              np.sqrt(maml._session.run(maml._loss, feed_dict=feed_dict))))

      loss2.append(

          np.average(

              np.sqrt(maml._session.run(maml._meta_loss, feed_dict=feed_dict))))

    # Initially the model should do a bad job of fitting the sine function.

    assert np.average(loss1) > 1.0

    # After one step of optimization it should do much better.

    assert np.average(loss2) < 1.0

    # Verify that we can create a new MAML object, reload the parameters from the first one, and

    # get the same result.

    new_maml = dc.metalearning.MAML(learner, model_dir=maml.model_dir)

    new_maml.restore()

    new_loss = np.average(np.sqrt(new_maml._session.run(new_maml._loss, feed_dict=feed_dict)))

    new_loss = np.average(

        np.sqrt(new_maml._session.run(new_maml._loss, feed_dict=feed_dict)))

    assert new_loss == loss1[-1]

    # Do the same thing, only using the "restore" argument to fit().

    new_maml = dc.metalearning.MAML(learner, model_dir=maml.model_dir)

    new_maml.fit(0, restore=True)

    new_loss = np.average(np.sqrt(new_maml._session.run(new_maml._loss, feed_dict=feed_dict)))

    new_loss = np.average(

        np.sqrt(new_maml._session.run(new_maml._loss, feed_dict=feed_dict)))

    assert new_loss == loss1[-1]