Commit 4bd621bc authored by galenxing's avatar galenxing
Browse files

upstream master layer

parent c0581e84
Loading
Loading
Loading
Loading
+0 −206
Original line number Diff line number Diff line
@@ -2614,209 +2614,3 @@ class AtomicConvolution(Layer):
    R = tf.reduce_sum(tf.multiply(D, D), 3)
    R = tf.sqrt(R)
    return R


def AlphaShare(in_layers=None, **kwargs):
  """
  This method should be used when constructing AlphaShare layers from Sluice Networks
  
  Parameters
  ----------
  in_layers: list of Layers or tensors
    tensors in list must be the same size and list must include two or more tensors

  Returns
  -------
  output_layers: list of Layers or tensors with same size as in_layers
    Distance matrix.

  References:
  Sluice networks: Learning what to share between loosely related tasks
  https://arxiv.org/abs/1705.08142
  """
  output_layers = []
  alpha_share = AlphaShareLayer(in_layers=in_layers, **kwargs)
  num_outputs = len(in_layers)
  for num_layer in range(0, num_outputs):
    ls = LayerSplitter(output_num=num_layer, in_layers=alpha_share)
    output_layers.append(ls)
  return output_layers


class AlphaShareLayer(Layer):
  """
  Part of a sluice network. Adds alpha parameters to control
  sharing between the main and auxillary tasks

  Factory method AlphaShare should be used for construction

  Parameters
  ----------
  in_layers: list of Layers or tensors
    tensors in list must be the same size and list must include two or more tensors

  Returns
  -------
  out_tensor: a tensor with shape [len(in_layers), x, y] where x, y were the original layer dimensions
    out_tensor should be fed into LayerSplitter 
  Distance matrix.
  """

  def __init__(self, **kwargs):
    super(AlphaShareLayer, self).__init__(**kwargs)

  def create_tensor(self, in_layers=None, set_tensors=True, **kwargs):
    inputs = self._get_input_tensors(in_layers)
    # check that there isnt just one or zero inputs
    if len(inputs) <= 1:
      raise ValueError("AlphaShare must have more than one input")
    self.num_outputs = len(inputs)
    # create subspaces
    subspaces = []
    original_cols = int(inputs[0].get_shape()[-1].value)
    subspace_size = int(original_cols / 2)
    for input_tensor in inputs:
      subspaces.append(tf.reshape(input_tensor[:, :subspace_size], [-1]))
      subspaces.append(tf.reshape(input_tensor[:, subspace_size:], [-1]))
    n_alphas = len(subspaces)
    subspaces = tf.reshape(tf.stack(subspaces), [n_alphas, -1])

    # create the alpha learnable parameters
    alphas = tf.Variable(tf.random_normal([n_alphas, n_alphas]), name='alphas')

    subspaces = tf.matmul(alphas, subspaces)

    # concatenate subspaces, reshape to size of original input, then stack
    # such that out_tensor has shape (2,?,original_cols)
    count = 0
    out_tensor = []
    tmp_tensor = []
    for row in range(n_alphas):
      tmp_tensor.append(tf.reshape(subspaces[row,], [-1, subspace_size]))
      count += 1
      if (count == 2):
        out_tensor.append(tf.concat(tmp_tensor, 1))
        tmp_tensor = []
        count = 0

    out_tensor = tf.stack(out_tensor)

    self.alphas = alphas
    if set_tensors:
      self.out_tensor = out_tensor
    return out_tensor

  def none_tensors(self):
    num_outputs, out_tensor, alphas = self.num_outputs, self.out_tensor, self.alphas
    self.num_outputs = None
    self.out_tensor = None
    self.alphas = None
    return num_outputs, out_tensor, alphas

  def set_tensors(self, tensor):
    self.num_outputs, self.out_tensor, self.alphas = tensor


class LayerSplitter(Layer):
  """
  Returns the nth output of a layer
  Assumes out_tensor has shape [x, :] where x is the total number of intended output tensors
  """

  def __init__(self, output_num, **kwargs):
    """
    Parameters
    ----------
    output_num: int
        returns the out_tensor[output_num, :] of a layer
    """
    self.output_num = output_num
    super(LayerSplitter, self).__init__(**kwargs)

  def create_tensor(self, in_layers=None, set_tensors=True, **kwargs):
    inputs = self._get_input_tensors(in_layers)[0]
    self.out_tensor = inputs[self.output_num, :]
    out_tensor = self.out_tensor
    return self.out_tensor

  def none_tensors(self):
    out_tensor = self.out_tensor
    self.out_tensor = None
    return out_tensor

  def set_tensors(self, tensor):
    self.out_tensor = tensor


class SluiceLoss(Layer):
  """
  Calculates the loss in a Sluice Network
  Every input into an AlphaShare should be used in SluiceLoss
  """

  def __init__(self, **kwargs):
    super(SluiceLoss, self).__init__(**kwargs)

  def create_tensor(self, in_layers=None, set_tensors=True, **kwargs):
    inputs = self._get_input_tensors(in_layers)
    temp = []
    subspaces = []
    # creates subspaces the same way it was done in AlphaShare
    for input_tensor in inputs:
      subspace_size = int(input_tensor.get_shape()[-1].value / 2)
      subspaces.append(input_tensor[:, :subspace_size])
      subspaces.append(input_tensor[:, subspace_size:])
      product = tf.matmul(tf.transpose(subspaces[0]), subspaces[1])
      subspaces = []
      # calculate squared Frobenius norm
      temp.append(tf.reduce_sum(tf.pow(product, 2)))
    out_tensor = tf.reduce_sum(temp)
    self.out_tensor = out_tensor
    return out_tensor


class BetaShare(Layer):
  """
  Part of a sluice network. Adds beta params to control which layer
  outputs are used for prediction

  Parameters
  ----------
  in_layers: list of Layers or tensors
    tensors in list must be the same size and list must include two or more tensors

  Returns
  -------
  output_layers: list of Layers or tensors with same size as in_layers
    Distance matrix.
  """

  def __init__(self, **kwargs):
    super(BetaShare, self).__init__(**kwargs)

  def create_tensor(self, in_layers=None, set_tensors=True, **kwargs):
    """
        Size of input layers must all be the same
        """
    inputs = self._get_input_tensors(in_layers)
    subspaces = []
    original_cols = int(inputs[0].get_shape()[-1].value)
    for input_tensor in inputs:
      subspaces.append(tf.reshape(input_tensor, [-1]))
    n_betas = len(inputs)
    subspaces = tf.reshape(tf.stack(subspaces), [n_betas, -1])

    betas = tf.Variable(tf.random_normal([1, n_betas]), name='betas')
    out_tensor = tf.matmul(betas, subspaces)
    self.betas = betas
    self.out_tensor = tf.reshape(out_tensor, [-1, original_cols])
    return self.out_tensor

  def none_tensors(self):
    out_tensor, betas = self.out_tensor, self.betas
    self.out_tensor = None
    self.betas = None
    return out_tensor, betas

  def set_tensors(self, tensor):
    self.out_tensor, self.betas = tensor
+2 −66
Original line number Diff line number Diff line
@@ -46,10 +46,6 @@ from deepchem.models.tensorgraph.layers import TensorWrapper
from deepchem.models.tensorgraph.layers import LSTMStep
from deepchem.models.tensorgraph.layers import AttnLSTMEmbedding
from deepchem.models.tensorgraph.layers import IterRefLSTMEmbedding
from deepchem.models.tensorgraph.layers import AlphaShareLayer
from deepchem.models.tensorgraph.layers import BetaShare
from deepchem.models.tensorgraph.layers import SluiceLoss
from deepchem.models.tensorgraph.layers import LayerSplitter

import deepchem as dc

@@ -130,7 +126,8 @@ class TestLayers(test_util.TensorFlowTestCase):
    dim = 2
    batch_size = 10
    mean_tensor = np.random.rand(dim)
    std_tensor = np.random.rand(1,)
    std_tensor = np.random.rand(
        1,)
    with self.test_session() as sess:
      mean_tensor = tf.convert_to_tensor(mean_tensor, dtype=tf.float32)
      std_tensor = tf.convert_to_tensor(std_tensor, dtype=tf.float32)
@@ -612,64 +609,3 @@ class TestLayers(test_util.TensorFlowTestCase):
      assert result == 1.5
      result = sess.run(tf.gradients(v, v))
      assert result[0] == 1.0

  def test_alpha_share_layer(self):
    """Test that alpha share works correctly"""
    batch_size = 50
    length = 10
    test_1 = np.random.rand(batch_size, length)
    test_2 = np.random.rand(batch_size, length)

    with self.test_session() as sess:
      test_1 = tf.convert_to_tensor(test_1, dtype=tf.float32)
      test_2 = tf.convert_to_tensor(test_2, dtype=tf.float32)

      out_tensor = AlphaShareLayer()(test_1, test_2)
      sess.run(tf.global_variables_initializer())
      test_1_out_tensor = out_tensor[0].eval()
      test_2_out_tensor = out_tensor[1].eval()
      assert test_1.shape == test_1_out_tensor.shape
      assert test_2.shape == test_2_out_tensor.shape

  def test_beta_share(self):
    """Test that beta share works correctly"""
    batch_size = 50
    length = 10
    test_1 = np.random.rand(batch_size, length)
    test_2 = np.random.rand(batch_size, length)

    with self.test_session() as sess:
      test_1 = tf.convert_to_tensor(test_1, dtype=tf.float32)
      test_2 = tf.convert_to_tensor(test_2, dtype=tf.float32)

      out_tensor = BetaShare()(test_1, test_2)
      sess.run(tf.global_variables_initializer())
      out_tensor.eval()
      assert test_1.shape == out_tensor.shape
      assert test_2.shape == out_tensor.shape

  def test_layer_splitter(self):
    """Test Layer Splitter"""
    input1 = np.arange(10).reshape(2, 5)
    input2 = np.arange(10, 20).reshape(2, 5)

    with self.test_session() as sess:
      input1 = tf.convert_to_tensor(input1, dtype=tf.float32)
      input2 = tf.convert_to_tensor(input2, dtype=tf.float32)
      input_tensor = tf.stack([input1, input2])
      output1 = LayerSplitter(0)(input_tensor)
      output2 = LayerSplitter(1)(input_tensor)
      sess.run(tf.global_variables_initializer())
      tf.assert_equal(input1, output1.eval())
      tf.assert_equal(input2, output2.eval())

  def test_sluice_loss(self):
    """Test the sluice loss function"""
    input1 = np.ones((3, 4))
    input2 = np.ones((2, 2))
    with self.test_session() as sess:
      input1 = tf.convert_to_tensor(input1, dtype=tf.float32)
      input2 = tf.convert_to_tensor(input2, dtype=tf.float32)
      output_tensor = SluiceLoss()(input1, input2)
      sess.run(tf.global_variables_initializer())
      assert output_tensor.eval() == 40.0