Merge pull request #1652 from peastman/callbacks

from deepchem.models.sklearn_models import SklearnModel

from deepchem.models.xgboost_models import XGBoostModel

from deepchem.models.multitask import SingletaskToMultitask

from deepchem.models.callbacks import ValidationCallback

from deepchem.models.tensorgraph.tensor_graph import TensorGraph

from deepchem.models.tensorgraph.fcnet import MultitaskRegressor

"""

Callback functions that can be invoked while fitting a KerasModel.

"""

from __future__ import print_function

import tensorflow as tf

import sys

class ValidationCallback(object):

  """Performs validation while training a KerasModel.

  This is a callback that can be passed to fit().  It periodically computes a

  set of metrics over a validation set and writes them to a file.  In addition,

  it can save the best model parameters found so far to a directory on disk,

  updating them every time it finds a new best validation score.

  If Tensorboard logging is enabled on the KerasModel, the metrics are also

  logged to Tensorboard.  This only happens when validation coincides with a

  step on which the model writes to the log.  You should therefore make sure

  that this callback's reporting interval is an even fraction or multiple of

  the model's logging interval.

  """

  def __init__(self,

               dataset,

               interval,

               metrics,

               output_file=sys.stdout,

               save_dir=None,

               save_metric=0,

               save_on_minimum=True):

    """Create a ValidationCallback.

    Parameters

    ----------

    dataset: dc.data.Dataset

      the validation set on which to compute the metrics

    interval: int

      the interval (in training steps) at which to perform validation

    metrics: list of dc.metrics.Metric

      metrics to compute on the validation set

    output_file: file

      to file to which results should be written

    save_dir: str

      if not None, the model parameters that produce the best validation score

      will be written to this directory

    save_metric: int

      the index of the metric to use when deciding whether to write a new set

      of parameters to disk

    save_on_minimum: bool

      if True, the best model is considered to be the one that minimizes the

      validation metric.  If False, the best model is considered to be the one

      that maximizes it.

    """

    self.dataset = dataset

    self.interval = interval

    self.metrics = metrics

    self.output_file = output_file

    self.save_dir = save_dir

    self.save_metric = save_metric

    self.save_on_minimum = save_on_minimum

    self._best_score = None

  def __call__(self, model, step):

    """This is invoked by the KerasModel after every step of fitting.

    Parameters

    ----------

    model: KerasModel

      the model that is being trained

    step: int

      the index of the training step that has just completed

    """

    if step % self.interval != 0:

      return

    scores = model.evaluate(self.dataset, self.metrics)

    message = 'Step %d validation:' % step

    for key in scores:

      message += ' %s=%g' % (key, scores[key])

    print(message, file=self.output_file)

    if model.tensorboard:

      for key in scores:

        model._log_value_to_tensorboard(tag=key, simple_value=scores[key])

    if self.save_dir is not None:

      score = scores[self.metrics[self.save_metric].name]

      if not self.save_on_minimum:

        score = -score

      if self._best_score is None or score < self._best_score:

        model.save_checkpoint(model_dir=self.save_dir)

        self._best_score = score

import time

import logging

import os

try:

  from collections.abc import Sequence

except:

  from collections import Sequence

logger = logging.getLogger(__name__)

      self.optimizer = optimizer

    self.tensorboard = tensorboard

    self.tensorboard_log_frequency = tensorboard_log_frequency

    self._tensorboard_step = 0

    if tensorboard and tf.executing_eagerly():

    if self.tensorboard:

      if tf.executing_eagerly():

        raise ValueError(

            "Logging to TensorBoard is not currently supported in eager mode")

      self._summary_writer = tf.summary.FileWriter(self.model_dir)

    if output_types is None:

      self._prediction_outputs = None

      self._loss_outputs = None

      # The loss doesn't depend on any variables.

      self._train_op = 0

    self._custom_train_op = {}

    if self.tensorboard:

      self._summary_ops = tf.summary.scalar('loss', self._loss_tensor)

      self._summary_writer = tf.summary.FileWriter(self.model_dir)

    self._init_new_vars()

  def _init_new_vars(self):

          deterministic=False,

          restore=False,

          variables=None,

          loss=None):

          loss=None,

          callbacks=[]):

    """Train this model on a dataset.

    Parameters

      a function of the form f(outputs, labels, weights) that computes the loss

      for each batch.  If None (the default), the model's standard loss function

      is used.

    callbacks: function or list of functions

      one or more functions of the form f(model, step) that will be invoked after

      every step.  This can be used to perform validation, logging, etc.

   """

    return self.fit_generator(

        self.default_generator(

            dataset, epochs=nb_epoch, deterministic=deterministic),

        max_checkpoints_to_keep, checkpoint_interval, restore, variables, loss)

            dataset, epochs=nb_epoch,

            deterministic=deterministic), max_checkpoints_to_keep,

        checkpoint_interval, restore, variables, loss, callbacks)

  def fit_generator(self,

                    generator,

                    checkpoint_interval=1000,

                    restore=False,

                    variables=None,

                    loss=None):

                    loss=None,

                    callbacks=[]):

    """Train this model on data from a generator.

    Parameters

      a function of the form f(outputs, labels, weights) that computes the loss

      for each batch.  If None (the default), the model's standard loss function

      is used.

    callbacks: function or list of functions

      one or more functions of the form f(model, step) that will be invoked after

      every step.  This can be used to perform validation, logging, etc.

    Returns

    -------

    the average loss over the most recent checkpoint interval

    """

    if not isinstance(callbacks, Sequence):

      callbacks = [callbacks]

    self._ensure_built()

    if checkpoint_interval > 0:

      manager = tf.train.CheckpointManager(self._checkpoint, self.model_dir,

        self.restore()

        restore = False

      inputs, labels, weights = self._prepare_batch(batch)

      self._tensorboard_step += 1

      should_log = (

          self._tensorboard_step % self.tensorboard_log_frequency == 0)

      self._current_summary = None

      if tf.executing_eagerly():

        # In eager mode we execute the loss function, accumulating the gradients.

          if self._loss_outputs is not None:

            outputs = [outputs[i] for i in self._loss_outputs]

          batch_loss = loss(outputs, labels, weights)

        avg_loss += batch_loss

        if variables is None:

          vars = self.model.trainable_variables

        else:

                  loss_tensor, global_step=self._global_step, var_list=vars)

            train_op = self._custom_train_op[op_key]

        fetches = [train_op, self._loss_tensor, self._global_step]

        if self.tensorboard and should_log:

          fetches.append(self._summary_ops)

        feed_dict = dict(zip(self._input_placeholders, inputs))

        feed_dict.update(dict(zip(self._label_placeholders, labels)))

        feed_dict.update(dict(zip(self._weights_placeholders, weights)))

        fetched_values = self.session.run(fetches, feed_dict=feed_dict)

        avg_loss += fetched_values[1]

        batch_loss = fetched_values[1]

        current_step = fetched_values[2]

        if self.tensorboard and should_log:

          self._summary_writer.reopen()

          self._summary_writer.add_summary(

              fetched_values[3], global_step=current_step)

          self._summary_writer.close()

      avg_loss += batch_loss

      # Report progress and write checkpoints.

      averaged_batches += 1

      should_log = (current_step % self.tensorboard_log_frequency == 0)

      if should_log:

        avg_loss = float(avg_loss) / averaged_batches

        logger.info(

      if checkpoint_interval > 0 and current_step % checkpoint_interval == checkpoint_interval - 1:

        self._exec_with_session(lambda: manager.save())

      for c in callbacks:

        c(self, current_step)

      if self.tensorboard and should_log:

        self._log_value_to_tensorboard(tag='loss', simple_value=batch_loss)

        self._summary_writer.reopen()

        self._summary_writer.add_summary(self._current_summary, current_step)

        self._summary_writer.close()

    # Report final results.

    if averaged_batches > 0:

    logger.info("TIMING: model fitting took %0.3f s" % (time2 - time1))

    return avg_loss

  def fit_on_batch(self, X, y, w, variables=None, loss=None):

  def _log_value_to_tensorboard(self, **kwargs):

    """This can be called during fitting to log a value to Tensorboard.

    Any keyword arguments passed to this method are passed on to summary.value.add().

    """

    if self._current_summary is None:

      self._current_summary = tf.Summary()

    self._current_summary.value.add(**kwargs)

  def fit_on_batch(self, X, y, w, variables=None, loss=None, callbacks=[]):

    """Perform a single step of training.

    Parameters

      a function of the form f(outputs, labels, weights) that computes the loss

      for each batch.  If None (the default), the model's standard loss function

      is used.

    callbacks: function or list of functions

      one or more functions of the form f(model, step) that will be invoked after

      every step.  This can be used to perform validation, logging, etc.

   """

    if not self.built:

      self.build()

    dataset = NumpyDataset(X, y, w)

    return self.fit(dataset, nb_epoch=1, variables=variables, loss=loss)

    return self.fit(

        dataset,

        nb_epoch=1,

        variables=variables,

        loss=loss,

        callbacks=callbacks)

  def _predict(self, generator, transformers, outputs, uncertainty):

    """

import unittest

import tempfile

import deepchem as dc

import numpy as np

import tensorflow as tf

from tensorflow.python.eager import context

try:

  from StringIO import StringIO

except ImportError:

  from io import StringIO

class TestKerasModel(unittest.TestCase):

  def test_validation(self):

    """Test ValidationCallback."""

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

    train_dataset, valid_dataset, test_dataset = datasets

    n_features = 1024

    model = dc.models.MultitaskClassifier(len(tasks), n_features, dropouts=0.5)

    # Train the model while logging the validation ROC AUC.

    metric = dc.metrics.Metric(dc.metrics.roc_auc_score, np.mean)

    log = StringIO()

    save_dir = tempfile.mkdtemp()

    callback = dc.models.ValidationCallback(

        valid_dataset,

        30, [metric],

        log,

        save_dir=save_dir,

        save_on_minimum=False)

    model.fit(train_dataset, callbacks=callback)

    # Parse the log to pull out the AUC scores.

    log.seek(0)

    scores = []

    for line in log:

      score = float(line.split('=')[-1])

      scores.append(score)

    # The last reported score should match the current performance of the model.

    valid_score = model.evaluate(valid_dataset, [metric], transformers)

    self.assertAlmostEqual(

        valid_score['mean-roc_auc_score'], scores[-1], places=5)

    # Reload the save model and confirm that it matches the best logged score.

    model.restore(model_dir=save_dir)

    valid_score = model.evaluate(valid_dataset, [metric], transformers)

    self.assertAlmostEqual(

        valid_score['mean-roc_auc_score'], max(scores), places=5)

  def test_validation_eager(self):

    """Test ValidationCallback, in eager mode."""

    with context.eager_mode():

      self.test_validation()

    """Test fitting a KerasModel defined as a graph."""

    n_data_points = 10

    n_features = 2

    np.random.seed(1234)

    X = np.random.rand(n_data_points, n_features)

    y = (X[:, 0] > X[:, 1]).astype(np.float32)

    dataset = dc.data.NumpyDataset(X, y)

      self.test_saliency_shapes()

  def test_tensorboard(self):

    """Test logging to Tensorboard."""

    n_data_points = 20

    n_features = 2

    X = np.random.rand(n_data_points, n_features)