Commit 21609311 authored by peastman's avatar peastman
Browse files

Losses support PyTorch

parent 84ab499f

deepchem/models/keras_model.py

+1 −1
Original line number Diff line number Diff line
@@ -1203,7 +1203,7 @@ class _StandardLoss(object): 
      raise ValueError(

          "Loss functions expects exactly one each of outputs, labels, and weights"

      )

    losses = self.loss(outputs[0], labels[0])

    losses = self.loss._compute_tf_loss(outputs[0], labels[0])

    w = weights[0]

    if len(w.shape) < len(losses.shape):

      if isinstance(w, tf.Tensor):

deepchem/models/losses.py

+79 −13
Original line number Diff line number Diff line 
import tensorflow as tf





class Loss:

  """A loss function for use in training models."""



  def __call__(self, output, labels):

    """Compute the loss function.

  def _compute_tf_loss(self, output, labels):

    """Compute the loss function for TensorFlow tensors.



    The inputs are tensors containing the model's outputs and the labels for a

    batch.  The return value should be a tensor of shape (batch_size) or
@@ -25,24 +22,38 @@ class Loss: 
    """

    raise NotImplementedError("Subclasses must implement this")



  def _create_pytorch_loss(self):

    """Create a PyTorch loss function."""

    raise NotImplementedError("Subclasses must implement this")





class L1Loss(Loss):

  """The absolute difference between the true and predicted values."""



  def __call__(self, output, labels):

  def _compute_tf_loss(self, output, labels):

    import tensorflow as tf

    output, labels = _make_shapes_consistent(output, labels)

    output, labels = _ensure_float(output, labels)

    return tf.abs(output - labels)



  def _create_pytorch_loss(self):

    import torch

    return torch.nn.L1Loss(reduction='none')





class L2Loss(Loss):

  """The squared difference between the true and predicted values."""



  def __call__(self, output, labels):

  def _compute_tf_loss(self, output, labels):

    import tensorflow as tf

    output, labels = _make_shapes_consistent(output, labels)

    output, labels = _ensure_float(output, labels)

    return tf.square(output - labels)



  def _create_pytorch_loss(self):

    import torch

    return torch.nn.MSELoss(reduction='none')





class HingeLoss(Loss):

  """The hinge loss function.
@@ -51,10 +62,19 @@ class HingeLoss(Loss): 
  should equal 0 or 1.

  """



  def __call__(self, output, labels):

  def _compute_tf_loss(self, output, labels):

    import tensorflow as tf

    output, labels = _make_shapes_consistent(output, labels)

    return tf.keras.losses.hinge(labels, output)



  def _create_pytorch_loss(self):

    import torch



    def loss(output, labels):

      return torch.mean(torch.clamp(1 - labels * output, min=0), dim=-1)



    return loss





class BinaryCrossEntropy(Loss):

  """The cross entropy between pairs of probabilities.
@@ -63,11 +83,21 @@ class BinaryCrossEntropy(Loss): 
  contain probabilities.

  """



  def __call__(self, output, labels):

  def _compute_tf_loss(self, output, labels):

    import tensorflow as tf

    output, labels = _make_shapes_consistent(output, labels)

    output, labels = _ensure_float(output, labels)

    return tf.keras.losses.binary_crossentropy(labels, output)



  def _create_pytorch_loss(self):

    import torch

    bce = torch.nn.BCELoss(reduction='none')



    def loss(output, labels):

      return torch.mean(bce(output, labels), dim=-1)



    return loss





class CategoricalCrossEntropy(Loss):

  """The cross entropy between two probability distributions.
@@ -77,11 +107,20 @@ class CategoricalCrossEntropy(Loss): 
  classes.

  """



  def __call__(self, output, labels):

  def _compute_tf_loss(self, output, labels):

    import tensorflow as tf

    output, labels = _make_shapes_consistent(output, labels)

    output, labels = _ensure_float(output, labels)

    return tf.keras.losses.categorical_crossentropy(labels, output)



  def _create_pytorch_loss(self):

    import torch



    def loss(output, labels):

      return -torch.sum(labels * torch.log(output), dim=-1)



    return loss





class SigmoidCrossEntropy(Loss):

  """The cross entropy between pairs of probabilities.
@@ -91,11 +130,21 @@ class SigmoidCrossEntropy(Loss): 
  converted to probabilities using a sigmoid function.

  """



  def __call__(self, output, labels):

  def _compute_tf_loss(self, output, labels):

    import tensorflow as tf

    output, labels = _make_shapes_consistent(output, labels)

    output, labels = _ensure_float(output, labels)

    return tf.nn.sigmoid_cross_entropy_with_logits(labels, output)



  def _create_pytorch_loss(self):

    import torch

    bce = torch.nn.BCELoss(reduction='none')



    def loss(output, labels):

      return bce(torch.sigmoid(output), labels)



    return loss





class SoftmaxCrossEntropy(Loss):

  """The cross entropy between two probability distributions.
@@ -106,11 +155,21 @@ class SoftmaxCrossEntropy(Loss): 
  function.

  """



  def __call__(self, output, labels):

  def _compute_tf_loss(self, output, labels):

    import tensorflow as tf

    output, labels = _make_shapes_consistent(output, labels)

    output, labels = _ensure_float(output, labels)

    return tf.nn.softmax_cross_entropy_with_logits(labels, output)



  def _create_pytorch_loss(self):

    import torch



    def loss(output, labels):

      return -torch.sum(

          labels * torch.log(torch.nn.functional.softmax(output, 1)), dim=-1)



    return loss





class SparseSoftmaxCrossEntropy(Loss):

  """The cross entropy between two probability distributions.
@@ -121,13 +180,19 @@ class SparseSoftmaxCrossEntropy(Loss): 
  using a softmax function.

  """



  def __call__(self, output, labels):

  def _compute_tf_loss(self, output, labels):

    import tensorflow as tf

    labels = tf.cast(labels, tf.int32)

    return tf.nn.sparse_softmax_cross_entropy_with_logits(labels, output)



  def _create_pytorch_loss(self):

    import torch

    return torch.nn.CrossEntropyLoss(reduction='none')





def _make_shapes_consistent(output, labels):

  """Try to make inputs have the same shape by adding dimensions of size 1."""

  import tensorflow as tf

  shape1 = output.shape

  shape2 = labels.shape

  len1 = len(shape1)
@@ -152,6 +217,7 @@ def _make_shapes_consistent(output, labels): 


def _ensure_float(output, labels):

  """Make sure the outputs and labels are both floating point types."""

  import tensorflow as tf

  if output.dtype not in (tf.float32, tf.float64):

    output = tf.cast(output, tf.float32)

  if labels.dtype not in (tf.float32, tf.float64):

deepchem/models/tests/test_losses.py

0 → 100644
+199 −0
Original line number Diff line number Diff line 
import deepchem.models.losses as losses

import unittest

import numpy as np



try:

  import tensorflow as tf

  has_tensorflow = True

except:

  has_tensorflow = False



try:

  import torch

  has_pytorch = True

except:

  has_pytorch = False





class TestLosses(unittest.TestCase):

  """Test loss functions."""



  @unittest.skipIf(not has_tensorflow, 'TensorFlow is not installed')

  def test_l1_loss_tf(self):

    """Test L1Loss."""

    loss = losses.L1Loss()

    outputs = tf.constant([[0.1, 0.8], [0.4, 0.6]])

    labels = tf.constant([[0.0, 1.0], [1.0, 0.0]])

    result = loss._compute_tf_loss(outputs, labels).numpy()

    expected = [[0.1, 0.2], [0.6, 0.6]]

    assert np.allclose(expected, result)



  @unittest.skipIf(not has_pytorch, 'PyTorch is not installed')

  def test_l1_loss_pytorch(self):

    """Test L1Loss."""

    loss = losses.L1Loss()

    outputs = torch.tensor([[0.1, 0.8], [0.4, 0.6]])

    labels = torch.tensor([[0.0, 1.0], [1.0, 0.0]])

    result = loss._create_pytorch_loss()(outputs, labels).numpy()

    expected = [[0.1, 0.2], [0.6, 0.6]]

    assert np.allclose(expected, result)



  @unittest.skipIf(not has_tensorflow, 'TensorFlow is not installed')

  def test_l2_loss_tf(self):

    """Test L2Loss."""

    loss = losses.L2Loss()

    outputs = tf.constant([[0.1, 0.8], [0.4, 0.6]])

    labels = tf.constant([[0.0, 1.0], [1.0, 0.0]])

    result = loss._compute_tf_loss(outputs, labels).numpy()

    expected = [[0.1**2, 0.2**2], [0.6**2, 0.6**2]]

    assert np.allclose(expected, result)



  @unittest.skipIf(not has_pytorch, 'PyTorch is not installed')

  def test_l2_loss_pytorch(self):

    """Test L2Loss."""

    loss = losses.L2Loss()

    outputs = torch.tensor([[0.1, 0.8], [0.4, 0.6]])

    labels = torch.tensor([[0.0, 1.0], [1.0, 0.0]])

    result = loss._create_pytorch_loss()(outputs, labels).numpy()

    expected = [[0.1**2, 0.2**2], [0.6**2, 0.6**2]]

    assert np.allclose(expected, result)



  @unittest.skipIf(not has_tensorflow, 'TensorFlow is not installed')

  def test_hinge_loss_tf(self):

    """Test HingeLoss."""

    loss = losses.HingeLoss()

    outputs = tf.constant([[0.1, 0.8], [0.4, 0.6]])

    labels = tf.constant([[1.0, -1.0], [-1.0, 1.0]])

    result = loss._compute_tf_loss(outputs, labels).numpy()

    expected = [np.mean([0.9, 1.8]), np.mean([1.4, 0.4])]

    assert np.allclose(expected, result)



  @unittest.skipIf(not has_pytorch, 'PyTorch is not installed')

  def test_hinge_loss_pytorch(self):

    """Test HingeLoss."""

    loss = losses.HingeLoss()

    outputs = torch.tensor([[0.1, 0.8], [0.4, 0.6]])

    labels = torch.tensor([[1.0, -1.0], [-1.0, 1.0]])

    result = loss._create_pytorch_loss()(outputs, labels).numpy()

    expected = [np.mean([0.9, 1.8]), np.mean([1.4, 0.4])]

    assert np.allclose(expected, result)



  @unittest.skipIf(not has_tensorflow, 'TensorFlow is not installed')

  def test_binary_cross_entropy_tf(self):

    """Test BinaryCrossEntropy."""

    loss = losses.BinaryCrossEntropy()

    outputs = tf.constant([[0.1, 0.8], [0.4, 0.6]])

    labels = torch.tensor([[0.0, 1.0], [1.0, 0.0]])

    result = loss._compute_tf_loss(outputs, labels).numpy()

    expected = [

        -np.mean([np.log(0.9), np.log(0.8)]),

        -np.mean([np.log(0.4), np.log(0.4)])

    ]

    assert np.allclose(expected, result)



  @unittest.skipIf(not has_pytorch, 'PyTorch is not installed')

  def test_binary_cross_entropy_pytorch(self):

    """Test BinaryCrossEntropy."""

    loss = losses.BinaryCrossEntropy()

    outputs = torch.tensor([[0.1, 0.8], [0.4, 0.6]])

    labels = torch.tensor([[0.0, 1.0], [1.0, 0.0]])

    result = loss._create_pytorch_loss()(outputs, labels).numpy()

    expected = [

        -np.mean([np.log(0.9), np.log(0.8)]),

        -np.mean([np.log(0.4), np.log(0.4)])

    ]

    assert np.allclose(expected, result)



  @unittest.skipIf(not has_tensorflow, 'TensorFlow is not installed')

  def test_categorical_cross_entropy_tf(self):

    """Test CategoricalCrossEntropy."""

    loss = losses.CategoricalCrossEntropy()

    outputs = tf.constant([[0.2, 0.8], [0.4, 0.6]])

    labels = torch.tensor([[0.0, 1.0], [1.0, 0.0]])

    result = loss._compute_tf_loss(outputs, labels).numpy()

    expected = [-np.log(0.8), -np.log(0.4)]

    assert np.allclose(expected, result)



  @unittest.skipIf(not has_pytorch, 'PyTorch is not installed')

  def test_categorical_cross_entropy_pytorch(self):

    """Test CategoricalCrossEntropy."""

    loss = losses.CategoricalCrossEntropy()

    outputs = torch.tensor([[0.2, 0.8], [0.4, 0.6]])

    labels = torch.tensor([[0.0, 1.0], [1.0, 0.0]])

    result = loss._create_pytorch_loss()(outputs, labels).numpy()

    expected = [-np.log(0.8), -np.log(0.4)]

    assert np.allclose(expected, result)



  @unittest.skipIf(not has_tensorflow, 'TensorFlow is not installed')

  def test_sigmoid_cross_entropy_tf(self):

    """Test SigmoidCrossEntropy."""

    loss = losses.SigmoidCrossEntropy()

    y = [[0.1, 0.8], [0.4, 0.6]]

    outputs = tf.constant(y)

    labels = torch.tensor([[0.0, 1.0], [1.0, 0.0]])

    result = loss._compute_tf_loss(outputs, labels).numpy()

    sigmoid = 1.0 / (1.0 + np.exp(-np.array(y)))

    expected = [[-np.log(1 - sigmoid[0, 0]), -np.log(sigmoid[0, 1])],

                [-np.log(sigmoid[1, 0]), -np.log(1 - sigmoid[1, 1])]]

    assert np.allclose(expected, result)



  @unittest.skipIf(not has_pytorch, 'PyTorch is not installed')

  def test_sigmoid_cross_entropy_pytorch(self):

    """Test SigmoidCrossEntropy."""

    loss = losses.SigmoidCrossEntropy()

    y = [[0.1, 0.8], [0.4, 0.6]]

    outputs = torch.tensor(y)

    labels = torch.tensor([[0.0, 1.0], [1.0, 0.0]])

    result = loss._create_pytorch_loss()(outputs, labels).numpy()

    sigmoid = 1.0 / (1.0 + np.exp(-np.array(y)))

    expected = [[-np.log(1 - sigmoid[0, 0]), -np.log(sigmoid[0, 1])],

                [-np.log(sigmoid[1, 0]), -np.log(1 - sigmoid[1, 1])]]

    assert np.allclose(expected, result)



  @unittest.skipIf(not has_tensorflow, 'TensorFlow is not installed')

  def test_softmax_cross_entropy_tf(self):

    """Test SoftmaxCrossEntropy."""

    loss = losses.SoftmaxCrossEntropy()

    y = np.array([[0.1, 0.8], [0.4, 0.6]])

    outputs = tf.constant(y)

    labels = torch.tensor([[0.0, 1.0], [1.0, 0.0]])

    result = loss._compute_tf_loss(outputs, labels).numpy()

    softmax = np.exp(y) / np.expand_dims(np.sum(np.exp(y), axis=1), 1)

    expected = [-np.log(softmax[0, 1]), -np.log(softmax[1, 0])]

    assert np.allclose(expected, result)



  @unittest.skipIf(not has_pytorch, 'PyTorch is not installed')

  def test_softmax_cross_entropy_pytorch(self):

    """Test SoftmaxCrossEntropy."""

    loss = losses.SoftmaxCrossEntropy()

    y = np.array([[0.1, 0.8], [0.4, 0.6]])

    outputs = torch.tensor(y)

    labels = torch.tensor([[0.0, 1.0], [1.0, 0.0]])

    result = loss._create_pytorch_loss()(outputs, labels).numpy()

    softmax = np.exp(y) / np.expand_dims(np.sum(np.exp(y), axis=1), 1)

    expected = [-np.log(softmax[0, 1]), -np.log(softmax[1, 0])]

    assert np.allclose(expected, result)



  @unittest.skipIf(not has_tensorflow, 'TensorFlow is not installed')

  def test_sparse_softmax_cross_entropy_tf(self):

    """Test SparseSoftmaxCrossEntropy."""

    loss = losses.SparseSoftmaxCrossEntropy()

    y = np.array([[0.1, 0.8], [0.4, 0.6]])

    outputs = tf.constant(y)

    labels = torch.tensor([1, 0])

    result = loss._compute_tf_loss(outputs, labels).numpy()

    softmax = np.exp(y) / np.expand_dims(np.sum(np.exp(y), axis=1), 1)

    expected = [-np.log(softmax[0, 1]), -np.log(softmax[1, 0])]

    assert np.allclose(expected, result)



  @unittest.skipIf(not has_pytorch, 'PyTorch is not installed')

  def test_sparse_softmax_cross_entropy_pytorch(self):

    """Test SparseSoftmaxCrossEntropy."""

    loss = losses.SparseSoftmaxCrossEntropy()

    y = np.array([[0.1, 0.8], [0.4, 0.6]])

    outputs = torch.tensor(y)

    labels = torch.tensor([1, 0])

    result = loss._create_pytorch_loss()(outputs, labels).numpy()

    softmax = np.exp(y) / np.expand_dims(np.sum(np.exp(y), axis=1), 1)

    expected = [-np.log(softmax[0, 1]), -np.log(softmax[1, 0])]

    assert np.allclose(expected, result)

deepchem/models/tests/test_optimizers.py

+1 −1
Original line number Diff line number Diff line
@@ -14,7 +14,7 @@ except: 
  has_pytorch = False





class TestLayers(unittest.TestCase):

class TestOptimizers(unittest.TestCase):

  """Test optimizers and related classes."""



  @unittest.skipIf(not has_tensorflow, 'TensorFlow is not installed')

CRA Git | Maintained and supported by SUSTech CRA and CCSE