|
|
| import torch |
| import torch.nn as nn |
| import torch.nn.functional as F |
|
|
| class StrongTextCNN(nn.Module): |
|
|
| def __init__( |
| self, |
| vocab_size, |
| embed_dim=300, |
| num_filters=128, |
| dropout=0.4 |
| ): |
| super().__init__() |
|
|
| self.embedding = nn.Embedding( |
| vocab_size, |
| embed_dim, |
| padding_idx=0 |
| ) |
|
|
| self.convs = nn.ModuleList([ |
| nn.Conv1d(embed_dim, num_filters, k) |
| for k in [2,3,4,5] |
| ]) |
|
|
| self.bn = nn.BatchNorm1d(num_filters * 4) |
|
|
| self.fc1 = nn.Linear(num_filters * 4, 256) |
| self.fc2 = nn.Linear(256, 64) |
| self.out = nn.Linear(64, 1) |
|
|
| self.dropout = nn.Dropout(dropout) |
|
|
| def forward(self, x): |
|
|
| x = self.embedding(x) |
| x = x.permute(0,2,1) |
|
|
| conv_outputs = [] |
|
|
| for conv in self.convs: |
|
|
| c = F.relu(conv(x)) |
|
|
| p = F.max_pool1d( |
| c, |
| kernel_size=c.shape[2] |
| ).squeeze(2) |
|
|
| conv_outputs.append(p) |
|
|
| x = torch.cat(conv_outputs, dim=1) |
|
|
| x = self.bn(x) |
|
|
| x = self.dropout(x) |
|
|
| x = F.relu(self.fc1(x)) |
| x = self.dropout(x) |
|
|
| x = F.relu(self.fc2(x)) |
| x = self.dropout(x) |
|
|
| x = self.out(x) |
|
|
| return x.squeeze(1) |
|
|