Commit 74360280 authored by Atreya Majumdar's avatar Atreya Majumdar
Browse files

Tests for encoder

parent 1aa112e2

deepchem/models/tests/test_layers.py

+68 −0
Original line number Diff line number Diff line
@@ -612,6 +612,7 @@ def test_scale_norm(): 
  """Test invoking ScaleNorm."""

  input_ar = torch.tensor([[1., 99., 10000.], [0.003, 999.37, 23.]])

  layer = torch_layers.ScaleNorm(0.35)

<<<<<<< HEAD

  result1 = layer(input_ar)

  output_ar = torch.tensor([[5.9157897e-05, 5.8566318e-03, 5.9157896e-01],

                            [1.7754727e-06, 5.9145141e-01, 1.3611957e-02]])
@@ -628,10 +629,58 @@ def test_multi_headed_mat_attention(): 
  adj_matrix = Chem.GetAdjacencyMatrix(mol)

  distance_matrix = Chem.GetDistanceMatrix(mol)

  layer = torch_layers.MultiHeadedMATAttention(

=======

  result1 = layer.forward(input_ar)

  output_ar = np.array([[5.9157897e-05, 5.8566318e-03, 5.9157896e-01],

                        [1.7754727e-06, 5.9145141e-01, 1.3611957e-02]])

  assert np.allclose(result1, output_ar)





@pytest.mark.torch

def test_position_wise_feed_forward():

  """Test invoking PositionwiseFeedForward."""

  torch.manual_seed(0)

  input_ar = torch.tensor([[1., 2.], [5., 6.]])

  layer = torch_layers.PositionwiseFeedForward(

      d_input=2,

      d_hidden=2,

      d_output=2,

      activation='relu',

      n_layers=1,

      dropout_p=0.0)

  result = layer(input_ar)

  output_ar = torch.tensor([[0.4810, 0.0000], [1.9771, 0.0000]])

  assert torch.allclose(result, output_ar, rtol=1e-4)





@pytest.mark.torch

def test_sub_layer_connection():

  """Test invoking SublayerConnection."""

  torch.manual_seed(0)

  input_ar = torch.tensor([[1., 2.], [5., 6.]])

  layer = torch_layers.SublayerConnection(2, 0.0)

  result = layer(input_ar, input_ar)

  output_ar = torch.tensor([[2.0027e-05, 3.0000e+00], [4.0000e+00, 7.0000e+00]])

  assert torch.allclose(result, output_ar)





@pytest.mark.torch

def test_mat_encoder_layer():

  """Test invoking MATEncoderLayer."""

  torch.manual_seed(0)

  import rdkit

  input_ar = torch.Tensor([[1., 2.], [5., 6.]])

  mask = torch.Tensor([[1., 1.], [1., 1.]])

  mol = rdkit.Chem.rdmolfiles.MolFromSmiles("CC")

  adj_matrix = rdkit.Chem.rdmolops.GetAdjacencyMatrix(mol)

  distance_matrix = rdkit.Chem.rdmolops.GetDistanceMatrix(mol)

  layer = torch_layers.MATEncoderLayer(

>>>>>>> Tests for encoder

      dist_kernel='softmax',

      lambda_attention=0.33,

      lambda_distance=0.33,

      h=2,

<<<<<<< HEAD

      hsize=2,

      dropout_p=0.0)

  input_tensor = torch.tensor([[1., 2.], [5., 6.]])
@@ -643,3 +692,22 @@ def test_multi_headed_mat_attention(): 
                            [[0.8671, 0.1069], [-3.4075, -0.8656],

                             [0.8671, 0.1069], [-3.4075, -0.8656]]])

  assert torch.allclose(result, output_ar, rtol=1e-3)

=======

      sa_hsize=2,

      sa_dropout_p=0.0,

      output_bias=True,

      d_input=2,

      d_hidden=2,

      d_output=2,

      activation='relu',

      n_layers=2,

      ff_dropout_p=0.0,

      encoder_hsize=2,

      encoder_dropout_p=0.0)

  result = layer(input_ar, mask, 0.0, adj_matrix, distance_matrix)

  output_ar = torch.tensor([[[0.9988, 2.0012], [-0.9999, 3.9999],

                             [0.9988, 2.0012], [-0.9999, 3.9999]],

                            [[5.0000, 6.0000], [3.0000, 8.0000],

                             [5.0000, 6.0000], [3.0000, 8.0000]]])

  assert torch.allclose(result, output_ar, rtol=1e-4)

>>>>>>> Tests for encoder

deepchem/models/torch_models/layers.py

+57 −20
Original line number Diff line number Diff line
@@ -66,6 +66,7 @@ class ScaleNorm(nn.Module): 
    return x * norm





<<<<<<< HEAD

<<<<<<< HEAD

class MultiHeadedMATAttention(nn.Module):

  """First constructs an attention layer tailored to the Molecular Attention Transformer [1]_ and then converts it into Multi-Headed Attention.
@@ -310,11 +311,13 @@ class MATEncoder(nn.Module): 
    return self.norm(x)





=======

>>>>>>> Tests for encoder

class MATEncoderLayer(nn.Module):

  """Encoder layer for use in the Molecular Attention Transformer [1]_.



  The MATEncoder layer is formed by adding self-attention and feed-forward to the encoder block.

  It is the basis of the MATEncoder block.

  The MATEncoder layer primarily consists of a self-attention layer (MultiHeadedMATAttention) and a feed-forward layer (PositionwiseFeedForward).

  This layer can be stacked multiple times to form an encoder.



  References

  ----------
@@ -323,7 +326,14 @@ class MATEncoderLayer(nn.Module): 
  Examples

  --------

  >>> import deepchem as dc

  >>> import rdkit

  >>> mol = rdkit.Chem.rdmolfiles.MolFromSmiles("CC")

  >>> adj_matrix = GetAdjacencyMatrix(mol)

  >>> distance_matrix = GetDistanceMatrix(mol)

  >>> layer = dc.models.torch_models.layers.MATEncoderLayer(dist_kernel = 'softmax', lambda_attention = 0.33, lambda_distance = 0.33, h = 8, sa_hsize = 1024, sa_dropout_p = 0.1, d_input = 1024, activation = 'relu', n_layers = 1, ff_dropout_p = 0.1, encoder_hsize = 1024, encoder_dropout_p = 0.1)

  >>> x = torch.Tensor([[1., 2.], [5., 6.]])

  >>> mask = torch.Tensor([[1., 1.], [1., 1.]])

  >>> output = layer(x, mask, sa_dropout_p = 0.0, adj_matrix = adj_matrix, distance_matrix = distance_matrix)

  """



  def __init__(self, dist_kernel: str, lambda_attention: float,
@@ -379,7 +389,14 @@ class MATEncoderLayer(nn.Module): 
    self.sublayer = nn.ModuleList([layer for _ in range(2)])

    self.size = encoder_hsize



  def forward(self, x: torch.Tensor, mask: torch.Tensor, **kwargs):

  def forward(

      self,

      x: torch.Tensor,

      mask: torch.Tensor,

      sa_dropout_p: float,

      adj_matrix: np.ndarray,

      distance_matrix: np.ndarray

  ):

    """Output computation for the MATEncoder layer.



    Parameters
@@ -388,10 +405,23 @@ class MATEncoderLayer(nn.Module): 
      Input tensor.

    mask: torch.Tensor

      Masks out padding values so that they are not taken into account when computing the attention score.

    sa_dropout_p: float

      Dropout probability for the self-attention layer (MultiHeadedMATAttention).

    adj_matrix: np.ndarray

      Adjacency matrix of a molecule.

    distance_matrix: np.ndarray

      Distance matrix of a molecule.

    """

    x = self.sublayer[0](x,

                         lambda x: self.self_attn(x, x, x, mask=mask, **kwargs))

    return self.sublayer[1](x, self.feed_forward)

                         self.self_attn(

                             x,

                             x,

                             x,

                             mask=mask,

                             dropout_p=sa_dropout_p,

                             adj_matrix=adj_matrix,

                             distance_matrix=distance_matrix))

    return self.sublayer[1](x, self.feed_forward(x))





class SublayerConnection(nn.Module):
@@ -409,7 +439,8 @@ class SublayerConnection(nn.Module): 
  >>> import deepchem as dc

  >>> scale = 0.35

  >>> layer = dc.models.torch_models.layers.SublayerConnection(2, 0.)

  >>> output = layer(torch.Tensor([1.,2.]), nn.Linear(2,1))

  >>> input_ar = torch.tensor([[1., 2.], [5., 6.]]) 

  >>> output = layer(input_ar, input_ar)

  """



  def __init__(self, size: int, dropout_p: float):
@@ -440,8 +471,20 @@ class SublayerConnection(nn.Module): 
      Layer whose normalized output will be added to x.

    """

    if x is None:

      return self.dropout(self.norm(output))

    return x + self.dropout(self.norm(output))

      return self.dropout_p(self.norm(output))



    if len(x.shape) < len(output.shape):

      temp_ar = x

      op_ar = output

      adjusted = temp_ar.unsqueeze(1).repeat(1, op_ar.shape[1], 1)

    elif len(x.shape) > len(output.shape):

      temp_ar = output

      op_ar = x

      adjusted = temp_ar.unsqueeze(1).repeat(1, op_ar.shape[1], 1)

    else:

      return x + self.dropout_p(self.norm(output))



    return adjusted + self.dropout_p(self.norm(op_ar))





class PositionwiseFeedForward(nn.Module):
@@ -460,13 +503,8 @@ class PositionwiseFeedForward(nn.Module): 
  >>> feed_fwd_layer = dc.models.torch_models.layers.PositionwiseFeedForward(d_input = 1024, d_hidden = None, d_output = None, activation = 'relu', n_layers = 1, dropout_p = 0.1)

  """



  def __init__(self,

               d_input: int,

               d_hidden: int,

               d_output: int,

               activation: str,

               n_layers: int,

               dropout_p: float):

  def __init__(self, d_input: int, d_hidden: int, d_output: int,

               activation: str, n_layers: int, dropout_p: float):

    """Initialize a PositionwiseFeedForward layer.



    Parameters
@@ -509,8 +547,8 @@ class PositionwiseFeedForward(nn.Module): 
      self.activation = lambda x: x



    self.n_layers = n_layers

    d_output = d_output if d_output is not 0 else d_input

    d_hidden = d_hidden if d_hidden is not 0 else d_input

    d_output = d_output if d_output != 0 else d_input

    d_hidden = d_hidden if d_hidden != 0 else d_input



    if n_layers == 1:

      self.linears = [nn.Linear(d_input, d_output)]
@@ -523,7 +561,6 @@ class PositionwiseFeedForward(nn.Module): 
    self.linears = nn.ModuleList(self.linears)

    dropout_layer = nn.Dropout(dropout_p)

    self.dropout_p = nn.ModuleList([dropout_layer for _ in range(n_layers)])

    self.act_func = activation



  def forward(self, x: torch.Tensor):

    """Output Computation for the PositionwiseFeedForward layer.
@@ -537,10 +574,10 @@ class PositionwiseFeedForward(nn.Module): 
      return x



    if self.n_layers == 1:

      return self.dropout_p[0](self.act_func(self.linears[0](x)))

      return self.dropout_p[0](self.activation(self.linears[0](x)))



    else:

      for i in range(self.n_layers - 1):

        x = self.dropout_p[i](self.act_func(self.linears[i](x)))

        x = self.dropout_p[i](self.activation(self.linears[i](x)))

      return self.linears[-1](x)

>>>>>>> Added encoder layers

docs/source/api_reference/layers.rst

+10 −0
Original line number Diff line number Diff line
@@ -125,7 +125,17 @@ Torch Layers 
.. autoclass:: deepchem.models.torch_models.layers.ScaleNorm

  :members:



<<<<<<< HEAD

.. autoclass:: deepchem.models.torch_models.layers.MultiHeadedMATAttention

=======

.. autoclass:: deepchem.models.torch_models.layers.MATEncoderLayer

  :members:



.. autoclass:: deepchem.models.torch_models.layers.SublayerConnection

  :members:



.. autoclass:: deepchem.models.torch_models.layers.PositionwiseFeedForward

>>>>>>> Tests for encoder

  :members:



.. autofunction:: deepchem.models.layers.cosine_dist

CRA Git | Maintained and supported by SUSTech CRA and CCSE