Initial implementation of MAML

from deepchem.metalearning.maml import MAML, MetaLearner

"""Model-Agnostic Meta-Learning (MAML) algorithm for low data learning."""

from deepchem.models.tensorgraph.layers import Layer

from deepchem.models.tensorgraph.optimizers import Adam

import os

import shutil

import tempfile

import tensorflow as tf

import time

class MetaLearner(object):

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

  To use MAML, create a subclass of this defining the learning problem to solve.

  It consists of a model that can be trained to perform many different tasks, and

  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."""

    raise NotImplemented("Subclasses must implement this")

  @property

  def variables(self):

    """Get the list of Tensorflow variables to train.

    The default implementation returns all trainable variables in the graph.  This is usually

    what you want, but subclasses can customize it if needed.

    """

    loss = self.loss

    if isinstance(loss, Layer):

        loss = loss.out_tensor

    with loss.graph.as_default():

        return tf.get_collection(tf.GraphKeys.TRAINABLE_VARIABLES)

  def select_task(self):

    """Select a new task to train on.

    If there is a fixed set of training tasks, this will typically cycle through them.

    If there are infinitely many training tasks, this can simply select a new one each

    time it is called.

    """

    raise NotImplemented("Subclasses must implement this")

  def get_batch(self):

    """Get a batch of data for training.

    This should return the data in the form of a Tensorflow feed dict, that is, a dict

    mapping tensors to values.  This will usually be called twice for each task, and should

    return a different batch on each call.

    """

    raise NotImplemented("Subclasses must implement this")

class MAML(object):

  """Implements the Model-Agnostic Meta-Learning algorithm for low data learning.

  The algorithm is described in Finn et al., "Model-Agnostic Meta-Learning for Fast

  Adaptation of Deep Networks" (https://arxiv.org/abs/1703.03400).  It is used for

  training models that can perform a variety of tasks, depending on what data they

  are trained on.  It assumes you have training data for many tasks, but only a small

  amount for each one.  It performs "meta-learning" by looping over tasks and trying

  to minimize the loss on each one *after* one or a few steps of gradient descent.

  That is, it does not try to create a model that can directly solve the tasks, but

  rather tries to create a model that is very easy to train.

  To use this class, create a subclass of MetaLearner that encapsulates the model

  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):

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

    Parameters

    ----------

    learner: MetaLearner

      defines the meta-learning problem

    learning_rate: float, Layer, or Tensor

      the learning rate to use for optimizing each task (not to be confused with the one used

      for meta-learning).  This can optionally be made a variable (represented as a Layer or

      Tensor), in which case the learning rate will itself be learnable.

    optimization_steps: int

      the number of steps of gradient descent to perform for each task

    optimizer: Optimizer

      the optimizer to use for meta-learning (not to be confused with the gradient descent

      optimization performed for each task)

    model_dir: str

      the directory in which the model will be saved.  If None, a temporary directory will be created.

    """

    # Record inputs.

    self.learner = learner

    if isinstance(learner.loss, Layer):

      self._loss = learner.loss.out_tensor

    else:

      self._loss = learner.loss

    if isinstance(learning_rate, Layer):

      self._learning_rate = learning_rate.out_tensor

    else:

      self._learning_rate = learning_rate

    self.optimizer = optimizer

    self._graph = self._loss.graph

    # Create the output directory if necessary.

    self._model_dir_is_temp = False

    if model_dir is not None:

      if not os.path.exists(model_dir):

        os.makedirs(model_dir)

    else:

      model_dir = tempfile.mkdtemp()

      self._model_dir_is_temp = True

    self.model_dir = model_dir

    self.save_file = "%s/%s" % (self.model_dir, "model")

    with self._graph.as_default():

      # Create duplicate placeholders for meta-optimization.

      learner.select_task()

      self._meta_placeholders = {}

      for p in learner.get_batch().keys():

        name = 'meta/'+p.name.split(':')[0]

        self._meta_placeholders[p] = tf.placeholder(p.dtype, p.shape, name)

      # Create the loss function for meta-optimization.

      updated_loss = self._loss

      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))

        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)

      self._meta_loss = updated_loss

      # Create the optimizer for meta-optimization.

      self._global_step = tf.placeholder(tf.int32, [])

      self._tf_optimizer = optimizer._create_optimizer(self._global_step)

      self._train_op = self._tf_optimizer.minimize(self._meta_loss)

      self._session = tf.Session()

  def __del__(self):

    if '_model_dir_is_temp' in dir(self) and self._model_dir_is_temp:

      shutil.rmtree(self.model_dir)

  def fit(self,

          steps,

          max_checkpoints_to_keep=5,

          checkpoint_interval=600,

          restore=False):

    """Perform meta-learning to train the model.

    Parameters

    ----------

    steps: int

      the number of steps of meta-learning to perform

    max_checkpoints_to_keep: int

      the maximum number of checkpoint files to keep.  When this number is reached, older

      files are deleted.

    checkpoint_interval: float

      the time interval at which to save checkpoints, measured in seconds

    restore: bool

      if True, restore the model from the most recent checkpoint and continue training

      from there.  If False, retrain the model from scratch.

    """

    with self._graph.as_default():

      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)

      checkpoint_index = 0

      checkpoint_time = time.time()

      # Main optimization loop.

      for i in range(steps):

        self.learner.select_task()

        feed_dict = self.learner.get_batch()

        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:

          saver.save(

              self._session,

              self.save_file,

              global_step=checkpoint_index)

          checkpoint_index += 1

          checkpoint_time = time.time()

  def restore(self):

    """Reload the model parameters from the most recent checkpoint file."""

    last_checkpoint = tf.train.latest_checkpoint(self.model_dir)

    if last_checkpoint is None:

      raise ValueError('No checkpoint found')

    with self._graph.as_default():

      saver = tf.train.Saver(self.learner.variables)

      saver.restore(self._session, last_checkpoint)

import deepchem as dc

from deepchem.models.tensorgraph.layers import Feature, Label, Dense, L2Loss

import numpy as np

import tensorflow as tf

import unittest

class TestMAML(unittest.TestCase):

  def test_sine(self):

    """Test meta-learning for sine function."""

    # This is a MetaLearner that learns to generate sine functions with variable

    # amplitude and phase.

    class SineLearner(dc.metalearning.MetaLearner):

      def __init__(self):

        self.batch_size = 10

        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)

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

        loss = 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()

      @property

      def loss(self):

        return self.tg.loss

      def select_task(self):

        self.amplitude = 5.0*np.random.random()

        self.phase = np.pi*np.random.random()

      def get_batch(self):

        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)

        return feed_dict

    # Optimize it.

    learner = SineLearner()

    maml = dc.metalearning.MAML(learner)

    maml.fit(25000)

    # Test it out on some new tasks and see how it works.

    loss1 = []

    loss2 = []

    for i in range(50):

      learner.select_task()

      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

    # 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)))

    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)))

    assert new_loss == loss1[-1]