| 489 | |
| 490 | |
| 491 | class PoswiseFeedForwardNet(layer.Layer): |
| 492 | def __init__(self, d_model=512, dim_feedforward=2048, bias=False): |
| 493 | super(PoswiseFeedForwardNet, self).__init__() |
| 494 | |
| 495 | self.d_model = d_model |
| 496 | self.dim_feedforward = dim_feedforward |
| 497 | self.bias = bias |
| 498 | |
| 499 | self.linear1 = Linear3D(d_model, dim_feedforward, bias=bias) |
| 500 | self.relu = layer.ReLU() |
| 501 | self.linear2 = Linear3D(dim_feedforward, d_model, bias=bias) |
| 502 | self.add = layer.Add() |
| 503 | self.norm = LayerNorm(d_model) |
| 504 | |
| 505 | def forward(self, inputs): |
| 506 | # inputs: [batch_size, seq_len, d_model] |
| 507 | residual = inputs |
| 508 | output = self.linear1(inputs) |
| 509 | output = self.relu(output) |
| 510 | output = self.linear2(output) |
| 511 | # [batch_size, seq_len, d_model] |
| 512 | output = self.add(output, residual) |
| 513 | output = self.norm(output) |
| 514 | return output |
| 515 | |
| 516 | |
| 517 | class LayerNorm(layer.Layer): |