1D convolutional text classifier in Pytorch.

1dcnn.py

import torch
import torch.nn as nn
import torch.nn.functional as fnn
import torch.optim as onn


class Transpose(nn.Module):

    def __init__(self):
        super(Transpose, self).__init__()

    def forward(self, tensor):
        return tensor.transpose(1, 2)


class Squeeze(nn.Module):

    def __init__(self):
        super(Squeeze, self).__init__()

    def forward(self, tensor):
        return tensor.squeeze()


def get_cnn(vocab_size, n_classes, embedding_size, kernel=2, frac=2):
    # assumes padded sequence length of 100 tokens.
    embeddings = nn.Embedding(
        vocab_size + 5, embedding_size,
        padding_idx=0,
    )
    encoder = nn.Sequential(
        nn.Conv1d(embedding_size, embedding_size // frac, kernel, padding=2),
        nn.ReLU(),
        nn.AvgPool1d(kernel_size=2),
        nn.Conv1d(embedding_size // frac, embedding_size // frac, kernel, padding=2),
        nn.ReLU(),
        nn.AvgPool1d(kernel_size=2),
        nn.Conv1d(embedding_size // frac, embedding_size // frac, 27),
        nn.ReLU(),
        Squeeze(),
        nn.Linear(embedding_size // frac, n_classes),
        nn.LogSoftmax(dim=1)
    )
    return nn.Sequential(
        embeddings,
        Transpose(),
        encoder
    )