Loading deepchem/models/tests/test_layers.py +3 −3 Original line number Diff line number Diff line Loading @@ -611,8 +611,8 @@ def test_DAG_gather(): def test_layer_norm(): """Test invoking LayerNorm.""" input_ar = torch.tensor([[1., 99., 10000.], [0.003, 999.37, 23.]]) layer = torch_layers.LayerNorm(input_ar.shape) layer = torch_layers.ScaleNorm(0.35) result1 = layer.forward(input_ar) output_ar = np.array([[-0.58585864, -0.5687999, 1.1546584], [-0.59738946, 1.1544659, -0.55707645]]) 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) deepchem/models/torch_models/layers.py +17 −25 Original line number Diff line number Diff line import math try: import torch import torch.nn as nn Loading @@ -5,40 +6,31 @@ except: raise ImportError('These classes require Torch to be installed.') class LayerNorm(nn.Module): """Apply Layer Normalization to input. class ScaleNorm(nn.Module): """Apply Scale Normalization to input. The layer takes input and applies layer normalization to it. All G values are initialized to sqrt(d). References ---------- .. [1] Lukasz Maziarka et al. "Molecule Attention Transformer" Graph Representation Learning workshop and Machine Learning and the Physical Sciences workshop at NeurIPS 2019. 2020. https://arxiv.org/abs/2002.08264 """ def __init__(self, features, eps=1e-6): """Initialize a LayerNorm layer. def __init__(self, scale, eps=1e-5): """Initialize a ScaleNorm layer. Parameters ---------- features: Tensor Tensor to be normalized. scale: Real number or single element tensor Scale magnitude. eps: float Epsilon value to be used. Epsilon value. """ super(LayerNorm, self).__init__() self.a_2 = nn.Parameter(torch.ones(features)) self.b_2 = nn.Parameter(torch.zeros(features)) super(ScaleNorm, self).__init__() self.scale = nn.Parameter(torch.tensor(math.sqrt(scale))) self.eps = eps def forward(self, x): """Normalize the Tensor. Parameters ---------- x: Tensor Tensor to be normalized. """ mean = x.mean(-1, keepdim=True) std = x.std(-1, keepdim=True) return self.a_2 * (x - mean) / (std + self.eps) + self.b_2 norm = self.scale / torch.norm(x, dim=-1, keepdim=True).clamp(min=self.eps) return x * norm docs/source/api_reference/layers.rst +1 −1 Original line number Diff line number Diff line Loading @@ -112,7 +112,7 @@ another tensor. DeepChem maintains an extensive collection of layers which perfo .. autoclass:: deepchem.models.layers.SetGather :members: .. autoclass:: deepchem.models.torch_models.layers.LayerNorm .. autoclass:: deepchem.models.torch_models.layers.ScaleNorm :members: .. autofunction:: deepchem.models.layers.cosine_dist Loading
deepchem/models/tests/test_layers.py +3 −3 Original line number Diff line number Diff line Loading @@ -611,8 +611,8 @@ def test_DAG_gather(): def test_layer_norm(): """Test invoking LayerNorm.""" input_ar = torch.tensor([[1., 99., 10000.], [0.003, 999.37, 23.]]) layer = torch_layers.LayerNorm(input_ar.shape) layer = torch_layers.ScaleNorm(0.35) result1 = layer.forward(input_ar) output_ar = np.array([[-0.58585864, -0.5687999, 1.1546584], [-0.59738946, 1.1544659, -0.55707645]]) 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)
deepchem/models/torch_models/layers.py +17 −25 Original line number Diff line number Diff line import math try: import torch import torch.nn as nn Loading @@ -5,40 +6,31 @@ except: raise ImportError('These classes require Torch to be installed.') class LayerNorm(nn.Module): """Apply Layer Normalization to input. class ScaleNorm(nn.Module): """Apply Scale Normalization to input. The layer takes input and applies layer normalization to it. All G values are initialized to sqrt(d). References ---------- .. [1] Lukasz Maziarka et al. "Molecule Attention Transformer" Graph Representation Learning workshop and Machine Learning and the Physical Sciences workshop at NeurIPS 2019. 2020. https://arxiv.org/abs/2002.08264 """ def __init__(self, features, eps=1e-6): """Initialize a LayerNorm layer. def __init__(self, scale, eps=1e-5): """Initialize a ScaleNorm layer. Parameters ---------- features: Tensor Tensor to be normalized. scale: Real number or single element tensor Scale magnitude. eps: float Epsilon value to be used. Epsilon value. """ super(LayerNorm, self).__init__() self.a_2 = nn.Parameter(torch.ones(features)) self.b_2 = nn.Parameter(torch.zeros(features)) super(ScaleNorm, self).__init__() self.scale = nn.Parameter(torch.tensor(math.sqrt(scale))) self.eps = eps def forward(self, x): """Normalize the Tensor. Parameters ---------- x: Tensor Tensor to be normalized. """ mean = x.mean(-1, keepdim=True) std = x.std(-1, keepdim=True) return self.a_2 * (x - mean) / (std + self.eps) + self.b_2 norm = self.scale / torch.norm(x, dim=-1, keepdim=True).clamp(min=self.eps) return x * norm
docs/source/api_reference/layers.rst +1 −1 Original line number Diff line number Diff line Loading @@ -112,7 +112,7 @@ another tensor. DeepChem maintains an extensive collection of layers which perfo .. autoclass:: deepchem.models.layers.SetGather :members: .. autoclass:: deepchem.models.torch_models.layers.LayerNorm .. autoclass:: deepchem.models.torch_models.layers.ScaleNorm :members: .. autofunction:: deepchem.models.layers.cosine_dist
