Commit 6200c834 authored by hsjang001205's avatar hsjang001205
Browse files

Add losses for VAE

parent 6d02180d

deepchem/models/losses.py

+69 −34
Original line number Diff line number Diff line
@@ -208,9 +208,34 @@ class SparseSoftmaxCrossEntropy(Loss): 
class VAE_ELBO(Loss):

  """The Variational AutoEncoder loss, KL Divergence Regularize + marginal log-likelihood.

  

  This losses basesd on "Auto-Encoding Variational Bayes" (https://arxiv.org/abs/1312.6114).

  ELBO(Evidence lower bound) lexically replaced Variational lower bound. 

  BCE means marginal log-likelihood, and KLD means KL divergence with normal distribution.

  Added hyper parameter 'kl_scale' for KLD.

  

  The logvar and mu should have shape (batch_size, hidden_space).

  The x and reconstruction_x should have (batch_size, attribute). 

  kl_scale: KLD regularized weights.

  The kl_scale should be float.

  

  Examples

  --------

  Examples for calculating loss using constant tensor.

  

  batch_size = 2,

  hidden_space = 2,

  num of original attribute = 3

  >>> logvar = np.array([[1.0,1.3],[0.6,1.2]])

  >>> mu = np.array([[0.2,0.7],[1.2,0.4]])

  >>> x = np.array([[0.9,0.4,0.8],[0.3,0,1]])

  >>> reconstruction = np.array([[0.8,0.3,0.7],[0.2,0,0.9]])

  

  Case tensorflow

  >>> VAE_ELBO()._compute_tf_loss(tf.constant(logvar), tf.constant(mu), tf.constant(x), tf.constant(reconstruction_x))

  <tf.Tensor: shape=(2,), dtype=float64, numpy=array([0.70165154, 0.76238271])>

  

  Case pytorch

  >>> (VAE_ELBO()._create_pytorch_loss())(torch.tensor(logvar), torch.tensor(mu), torch.tensor(x), torch.tensor(reconstruction_x))

  tensor([0.7017, 0.7624], dtype=torch.float64)

  """



  def _compute_tf_loss(self, logvar, mu, x, reconstruction_x, kl_scale = 1):
@@ -227,7 +252,7 @@ class VAE_ELBO(Loss): 


    def loss(logvar, mu, x, reconstruction_x, kl_scale = 1):

      x, reconstruction_x = _make_pytorch_shapes_consistent(x, reconstruction_x)

      BCE = torch.mean(bce(x, reconstruction_x), dim=-1)

      BCE = torch.mean(bce(reconstruction_x, x), dim=-1)

      KLD = (VAE_KLDivergence()._create_pytorch_loss())(logvar, mu)

      return BCE + kl_scale*KLD
@@ -235,10 +260,31 @@ class VAE_ELBO(Loss): 




class VAE_KLDivergence(Loss):

  """The KL_divergence between hidden distribution and normal distribution

  """The KL_divergence between hidden distribution and normal distribution.

  

  This loss implements KL divergence losses with normal distribution 

  based on "Auto-Encoding Variational Bayes" (https://arxiv.org/abs/1312.6114).

  

  The logvar should have shape (batch_size, hidden_space) and each term represents

  standard deviation of hidden distribution. The mean shuold have 

  (batch_size, hidden_space) and each term represents mean of hidden distribtuon.

  

  Examples

  --------

  Examples for calculating loss using constant tensor.

  

  batch_size = 2,

  hidden_space = 2,

  >>> logvar = np.array([[1.0,1.3],[0.6,1.2]])

  >>> mu = np.array([[0.2,0.7],[1.2,0.4]])

  

  Case tensorflow

  >>> VAE_KLDivergence()._compute_tf_loss(tf.constant(logvar), tf.constant(mu))

  tf.Tensor([0.52783368 0.24813068], shape=(2,), dtype=float64)

  

  Case pytorch

  >>> (VAE_KLDivergence()._create_pytorch_loss())(torch.tensor(logvar), torch.tensor(mu))

  tensor([0.1738, 0.5143], dtype=torch.float64)

  """



  def _compute_tf_loss(self, logvar, mu):
@@ -257,41 +303,30 @@ class VAE_KLDivergence(Loss): 
    return loss





class KLDivergence(Loss):

  """The KL_divergence between two distribution D_KL(P||Q).

  The argument should have shape (batch_size, num of variable) and represents

  probabilites distribution. 

  """



  def _compute_tf_loss(self, P, Q):

    import tensorflow as tf

    P, Q = _make_tf_shapes_consistent(P, Q)

    P, Q = _ensure_float(P, Q)

    #extended one of probabilites to to binary distribution

    if P.shape[-1] == 1:

      P = tf.concat([P,1-P], axis = -1)

      Q = tf.concat([Q,1-Q], axis = -1)

    return tf.reduce_mean(P * tf.math.log((P+1e-20) / (Q+1e-20)), axis=-1)

class ShannonEntropy(Loss):

  """The ShannonEntropy of discrete-distribution.

  

  def _create_pytorch_loss(self):

    import torch

  This loss implements shannon entropy based on

  "A Brief Introduction to Shannon's Information Theory" (https://arxiv.org/abs/1612.09316).

  

    def loss(P, Q):

      P, Q = _make_pytorch_shapes_consistent(P, Q)

      #extended one of probabilites to binary distribution

      if P.shape[-1] == 1:

        P = torch.cat((P,1-P), dim = -1)

        Q = torch.cat((Q,1-Q), dim = -1)

      return torch.mean(P * torch.log((P+1e-20) / (Q+1e-20)),dim = -1)

  The inputs should have shape (batch size, num of variable) and represents

  probabilites distribution.

  

    return loss

  Examples

  --------

  Examples for calculating loss using constant tensor.

  

  batch_size = 2,

  num_of variable = variable,

  >>> inputs = np.array([[0.7,0.3],[0.9,0.1]])

  

class ShannonEntropy(Loss):

  """The ShannonEntropy of discrete-distribution.

  Case tensorflow

  >>> ShannonEntropy()._compute_tf_loss(tf.constant(inputs))

  tf.Tensor([0.52783368 0.24813068], shape=(2,), dtype=float64)

  

  The inputs should have shape (batch size, num of variable) and represents

  probabilites distribution.

  Case pytorch

  >>> (ShannonEntropy()._create_pytorch_loss())(torch.tensor(inputs))

  tensor([0.1738, 0.5143], dtype=torch.float64)

  """



  def _compute_tf_loss(self, inputs):

deepchem/models/tests/test_losses.py

+76 −1
Original line number Diff line number Diff line 
import deepchem.models.losses as losses

#import deepchem.models.losses as losses

import sys

sys.path.append('C:/Users/hsjang/aa/deepchem/deepchem/models')

import losses

import unittest

import numpy as np
@@ -197,3 +200,75 @@ class TestLosses(unittest.TestCase): 
    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_VAE_ELBO_tf(self):

    """."""

    loss = losses.VAE_ELBO()

    logvar = tf.constant([[1.0,1.3],[0.6,1.2]])

    mu = tf.constant([[0.2,0.7],[1.2,0.4]])

    x = tf.constant([[0.9,0.4,0.8],[0.3,0,1]])

    reconstruction_x = tf.constant([[0.8,0.3,0.7],[0.2,0,0.9]])

    result = loss._compute_tf_loss(logvar, mu, x, reconstruction_x).numpy()

    expected = [0.5 * np.mean([0.04+1.0-np.log(1e-20+1.0)-1, 0.49+1.69 - np.log(1e-20 +1.69) - 1])

                -np.mean(np.array([0.9,0.4,0.8])*np.log([0.8,0.3,0.7])+np.array([0.1,0.6,0.2])*np.log([0.2,0.7,0.3])),

                0.5 * np.mean([1.44+0.36-np.log(1e-20+0.36)-1, 0.16+1.44 - np.log(1e-20 +1.44) - 1])

                -np.mean(np.array([0.3,0,1])*np.log([0.2,1e-20,0.9])+np.array([0.7,1,0])*np.log([0.8,1,0.1]))]

    assert np.allclose(expected, result)

    

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

  def test_VAE_ELBO_pytorch(self):

    """."""

    loss = losses.VAE_ELBO()

    logvar = torch.tensor([[1.0,1.3],[0.6,1.2]])

    mu = torch.tensor([[0.2,0.7],[1.2,0.4]])

    x = torch.tensor([[0.9,0.4,0.8],[0.3,0,1]])

    reconstruction_x = torch.tensor([[0.8,0.3,0.7],[0.2,0,0.9]])

    result = loss._create_pytorch_loss()(logvar, mu, x, reconstruction_x).numpy()

    expected = [0.5 * np.mean([0.04+1.0-np.log(1e-20+1.0)-1, 0.49+1.69 - np.log(1e-20 +1.69) - 1])

                -np.mean(np.array([0.9,0.4,0.8])*np.log([0.8,0.3,0.7])+np.array([0.1,0.6,0.2])*np.log([0.2,0.7,0.3])),

                0.5 * np.mean([1.44+0.36-np.log(1e-20+0.36)-1, 0.16+1.44 - np.log(1e-20 +1.44) - 1])

                -np.mean(np.array([0.3,0,1])*np.log([0.2,1e-20,0.9])+np.array([0.7,1,0])*np.log([0.8,1,0.1]))]

    assert np.allclose(expected, result)



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

  def test_VAE_KLDivergence_tf(self):

    """."""

    loss = losses.VAE_KLDivergence()

    logvar = tf.constant([[1.0,1.3],[0.6,1.2]])

    mu = tf.constant([[0.2,0.7],[1.2,0.4]])

    result = loss._compute_tf_loss(logvar, mu).numpy()

    expected = [0.5 * np.mean([0.04+1.0-np.log(1e-20+1.0)-1, 0.49+1.69 - np.log(1e-20 +1.69) - 1]),

                0.5 * np.mean([1.44+0.36-np.log(1e-20+0.36)-1, 0.16+1.44 - np.log(1e-20 +1.44) - 1])]

    assert np.allclose(expected, result)



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

  def test_VAE_KLDivergence_pytorch(self):

    """."""

    loss = losses.VAE_KLDivergence()

    logvar = torch.tensor([[1.0,1.3],[0.6,1.2]])

    mu = torch.tensor([[0.2,0.7],[1.2,0.4]])

    result = loss._create_pytorch_loss()(logvar, mu).numpy()

    expected = [0.5 * np.mean([0.04+1.0-np.log(1e-20+1.0)-1, 0.49+1.69 - np.log(1e-20 +1.69) - 1]),

                0.5 * np.mean([1.44+0.36-np.log(1e-20+0.36)-1, 0.16+1.44 - np.log(1e-20 +1.44) - 1])]

    assert np.allclose(expected, result)

    

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

  def test_ShannonEntropy_tf(self):

    """."""

    loss = losses.ShannonEntropy()

    inputs = tf.constant([[0.7,0.3],[0.9,0.1]])

    result = loss._compute_tf_loss(inputs).numpy()

    expected = [-np.mean([0.7*np.log(0.7),0.3*np.log(0.3)]),

                -np.mean([0.9*np.log(0.9),0.1*np.log(0.1)])]

    assert np.allclose(expected, result)



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

  def test_ShannonEntropy_pytorch(self):

    """."""

    loss = losses.ShannonEntropy()

    inputs = torch.tensor([[0.7,0.3],[0.9,0.1]])

    result = loss._create_pytorch_loss()(inputs).numpy()

    expected = [-np.mean([0.7*np.log(0.7),0.3*np.log(0.3)]),

                -np.mean([0.9*np.log(0.9),0.1*np.log(0.1)])]

    assert np.allclose(expected, result)
 
 No newline at end of file

CRA Git | Maintained and supported by SUSTech CRA and CCSE