Merge pull request #759 from peastman/maml

import deepchem.data

import deepchem.feat

import deepchem.hyper

import deepchem.metalearning

import deepchem.metrics

import deepchem.models

import deepchem.nn

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, GradientDescent

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.

  """

  @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().

  You can then use train_on_current_task() to fine tune the model for a particular

  task.

  """

  def __init__(self,

               learner,

               learning_rate=0.001,

               optimization_steps=1,

               meta_batch_size=10,

               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

    meta_batch_size: int

      the number of tasks to use for each step of meta-learning

    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.meta_batch_size = meta_batch_size

    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 variables for accumulating the gradients.

      gradients = tf.gradients(self._meta_loss, learner.variables)

      zero_gradients = [tf.zeros(g.shape, g.dtype) for g in gradients]

      summed_gradients = [

          tf.Variable(z, trainable=False) for z in zero_gradients

      ]

      self._clear_gradients = tf.group(

          * [s.assign(z) for s, z in zip(summed_gradients, zero_gradients)])

      self._add_gradients = tf.group(

          * [s.assign_add(g) for s, g in zip(summed_gradients, gradients)])

      # Create the optimizers for meta-optimization and task optimization.

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

      grads_and_vars = list(zip(summed_gradients, learner.variables))

      self._meta_train_op = optimizer._create_optimizer(

          self._global_step).apply_gradients(grads_and_vars)

      task_optimizer = GradientDescent(learning_rate=self._learning_rate)

      self._task_train_op = task_optimizer._create_optimizer(

          self._global_step).minimize(self._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._session.run(self._clear_gradients)

        for j in range(self.meta_batch_size):

          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

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

        self._session.run(self._meta_train_op)

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

  def train_on_current_task(self, optimization_steps=1, restore=True):

    """Perform a few steps of gradient descent to fine tune the model on the current task.

    Parameters

    ----------

    optimization_steps: int

      the number of steps of gradient descent to perform

    restore: bool

      if True, restore the model from the most recent checkpoint before optimizing

    """

    if restore:

      self.restore()

    with self._graph.as_default():

      feed_dict = self.learner.get_batch()

      for i in range(optimization_steps):

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

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, meta_batch_size=4)

    maml.fit(12000)

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

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

    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]

from __future__ import print_function

import deepchem as dc

import numpy as np

import tensorflow as tf

from sklearn.metrics import accuracy_score

# Load the data.

tasks, datasets, transformers = dc.molnet.load_toxcast()

(train_dataset, valid_dataset, test_dataset) = datasets

x = train_dataset.X

y = train_dataset.y

w = train_dataset.w

n_features = x.shape[1]

n_molecules = y.shape[0]

n_tasks = y.shape[1]

# Toxcast has data on 6874 molecules and 617 tasks.  However, the data is very

# sparse: most tasks do not include data for most molecules.  It also is very

# unbalanced: there are many more negatives than positives.  For each task,

# create a list of alternating postives and negatives so each batch will have

# equal numbers of both.

task_molecules = []

for i in range(n_tasks):

  positives = [j for j in range(n_molecules) if w[j, i] > 0 and y[j, i] == 1]

  negatives = [j for j in range(n_molecules) if w[j, i] > 0 and y[j, i] == 0]

  np.random.shuffle(positives)

  np.random.shuffle(negatives)

  mols = sum((list(x) for x in zip(positives, negatives)), [])

  task_molecules.append(mols)

# Create the model to train.  We use a simple fully connected network with

# one hidden layer.

model = dc.models.TensorGraphMultiTaskClassifier(

    1, n_features, layer_sizes=[1000], dropouts=[0.0])

model.build()

# Define a MetaLearner describing the learning problem.

class ToxcastLearner(dc.metalearning.MetaLearner):

  def __init__(self):

    self.n_training_tasks = int(n_tasks * 0.8)

    self.batch_size = 10

    self.batch_start = [0] * n_tasks

    self.set_task_index(0)

  @property

  def loss(self):

    return model.loss

  def set_task_index(self, index):

    self.task = index

  def select_task(self):

    self.set_task_index((self.task + 1) % self.n_training_tasks)

  def get_batch(self):

    task = self.task

    start = self.batch_start[task]

    mols = task_molecules[task][start:start + self.batch_size]

    labels = np.zeros((self.batch_size, 1, 2))

    labels[np.arange(self.batch_size), 0, y[mols, task].astype(np.int64)] = 1

    weights = np.ones((self.batch_size, 1))

    feed_dict = {}

    feed_dict[model.features[0].out_tensor] = x[mols, :]

    feed_dict[model.labels[0].out_tensor] = labels

    feed_dict[model.task_weights[0].out_tensor] = weights

    if start + 2 * self.batch_size > len(task_molecules[task]):

      self.batch_start[task] = 0

    else:

      self.batch_start[task] += self.batch_size

    return feed_dict

# Run meta-learning on 80% of the tasks.

n_epochs = 20

learner = ToxcastLearner()

maml = dc.metalearning.MAML(learner)

steps = n_epochs * learner.n_training_tasks // maml.meta_batch_size

maml.fit(steps)

# Validate on the remaining tasks.

def compute_scores(optimize):

  maml.restore()

  y_true = []

  y_pred = []

  losses = []

  with model._get_tf("Graph").as_default():

    prediction = tf.contrib.layers.softmax(model.outputs[0].out_tensor)

    for task in range(learner.n_training_tasks, n_tasks):

      learner.set_task_index(task)

      if optimize:

        maml.train_on_current_task()

      feed_dict = learner.get_batch()

      y_true.append(feed_dict[model.labels[0].out_tensor][:, 0, 0])

      y_pred.append(maml._session.run(prediction, feed_dict=feed_dict)[:, 0, 0])

      losses.append(maml._session.run(model.loss, feed_dict=feed_dict))

  y_true = np.concatenate(y_true)

  y_pred = np.concatenate(y_pred)

  print()

  print('Cross entropy loss:', np.mean(losses))

  print('Prediction accuracy:', accuracy_score(y_true, y_pred > 0.5))

  print('ROC AUC:', dc.metrics.compute_roc_auc_scores(y_true, y_pred))

  print()

print('Before fine tuning:')

compute_scores(False)

print('After fine tuning:')

compute_scores(True)