| 271 | print("Better speed can be achieved with apex installed from https://www.github.com/nvidia/apex.") |
| 272 | |
| 273 | class BertLayerNorm(nn.Module): |
| 274 | |
| 275 | def __init__(self, hidden_size, eps=1e-12): |
| 276 | """Construct a layernorm module in the TF style (epsilon inside the square root). |
| 277 | """ |
| 278 | super(BertLayerNorm, self).__init__() |
| 279 | self.weight = nn.Parameter(torch.ones(hidden_size)) |
| 280 | self.bias = nn.Parameter(torch.zeros(hidden_size)) |
| 281 | self.variance_epsilon = eps |
| 282 | |
| 283 | def forward(self, x): |
| 284 | u = x.mean(-1, keepdim=True) |
| 285 | s = (x - u).pow(2).mean(-1, keepdim=True) |
| 286 | x = (x - u) / torch.sqrt(s + self.variance_epsilon) |
| 287 | return self.weight * x + self.bias |
| 288 | |
| 289 | |
| 290 | class BertSelfAttention(nn.Module): |