[XLNet] Use pytorch's layernorm like in BERT

See huggingface#1089 cc @thomwolf @LysandreJik Also @dhpollack
sorami · Aug 31, 2019 · 1d438f1 · 1d438f1
1 parent 574c5b3
commit 1d438f1
Showing 1 changed file with 1 addition and 14 deletions.
diff --git a/pytorch_transformers/modeling_xlnet.py b/pytorch_transformers/modeling_xlnet.py
@@ -337,20 +337,7 @@ def num_hidden_layers(self):
     from apex.normalization.fused_layer_norm import FusedLayerNorm as XLNetLayerNorm
 except (ImportError, AttributeError) as e:
     logger.info("Better speed can be achieved with apex installed from https://www.github.com/nvidia/apex .")
-    class XLNetLayerNorm(nn.Module):
-        def __init__(self, d_model, eps=1e-12):
-            """Construct a layernorm module in the TF style (epsilon inside the square root).
-            """
-            super(XLNetLayerNorm, self).__init__()
-            self.weight = nn.Parameter(torch.ones(d_model))
-            self.bias = nn.Parameter(torch.zeros(d_model))
-            self.variance_epsilon = eps
-
-        def forward(self, x):
-            u = x.mean(-1, keepdim=True)
-            s = (x - u).pow(2).mean(-1, keepdim=True)
-            x = (x - u) / torch.sqrt(s + self.variance_epsilon)
-            return self.weight * x + self.bias
+    from torch.nn import LayerNorm as XLNetLayerNorm
 
 class XLNetRelativeAttention(nn.Module):
     def __init__(self, config):