| """ | |
| Implementation of "Convolutional Sequence to Sequence Learning" | |
| """ | |
| import torch | |
| import torch.nn as nn | |
| import torch.nn.init as init | |
| import onmt.modules | |
| SCALE_WEIGHT = 0.5 ** 0.5 | |
| def shape_transform(x): | |
| """ Tranform the size of the tensors to fit for conv input. """ | |
| return torch.unsqueeze(torch.transpose(x, 1, 2), 3) | |
| class GatedConv(nn.Module): | |
| """ Gated convolution for CNN class """ | |
| def __init__(self, input_size, width=3, dropout=0.2, nopad=False): | |
| super(GatedConv, self).__init__() | |
| self.conv = onmt.modules.WeightNormConv2d( | |
| input_size, 2 * input_size, kernel_size=(width, 1), stride=(1, 1), | |
| padding=(width // 2 * (1 - nopad), 0)) | |
| init.xavier_uniform_(self.conv.weight, gain=(4 * (1 - dropout))**0.5) | |
| self.dropout = nn.Dropout(dropout) | |
| def forward(self, x_var): | |
| x_var = self.dropout(x_var) | |
| x_var = self.conv(x_var) | |
| out, gate = x_var.split(int(x_var.size(1) / 2), 1) | |
| out = out * torch.sigmoid(gate) | |
| return out | |
| class StackedCNN(nn.Module): | |
| """ Stacked CNN class """ | |
| def __init__(self, num_layers, input_size, cnn_kernel_width=3, | |
| dropout=0.2): | |
| super(StackedCNN, self).__init__() | |
| self.dropout = dropout | |
| self.num_layers = num_layers | |
| self.layers = nn.ModuleList() | |
| for _ in range(num_layers): | |
| self.layers.append( | |
| GatedConv(input_size, cnn_kernel_width, dropout)) | |
| def forward(self, x): | |
| for conv in self.layers: | |
| x = x + conv(x) | |
| x *= SCALE_WEIGHT | |
| return x | |