Training a Document Classification Model in PyTorch with Hierarchical Attention

Document classification is a fundamental task in Natural Language Processing (NLP), where the goal is to categorize texts into predefined classes. With the rise of deep learning, models leveraging attention mechanisms, particularly the Hierarchical Attention Network (HAN), have gained popularity due to their ability to focus on important parts of a document, mimicking human-like reading habits.

In this article, we'll guide you step-by-step to train a document classification model using the Hierarchical Attention Model in PyTorch. This approach effectively processes document structures, considering word sequences and sentence hierarchical information.

Prerequisites

Before we start, ensure you have the following libraries installed:

pip install torch torchtext

Additionally, some familiarity with PyTorch and basic NLP concepts is recommended.

Data Preparation

Let's say we are working with news articles, and each article belongs to categories like Sports, Technology, Politics, etc. You will need data where each document is labeled with its category. For data processing, we can use torchtext for tokenization and batch generation.

from torchtext.datasets import AG_NEWS
from torchtext.data.utils import get_tokenizer

train_iter, test_iter = AG_NEWS(split=('train', 'test'))
tokenizer = get_tokenizer('basic_english')

Convert your text into tensors using the vocabularies built from your tokenizer.

Model Structure

The Hierarchical Attention Model consists of several layers: a word encoder, a word-level attention layer, a sentence encoder, and a sentence-level attention layer.

1. Word Encoder

This layer embeds the words using an embedding layer or pre-trained embeddings such as GloVe.

from torch import nn

class WordEncoder(nn.Module):
    def __init__(self, vocab_size, embed_size):
        super(WordEncoder, self).__init__()
        self.embedding = nn.Embedding(vocab_size, embed_size)

    def forward(self, x):
        return self.embedding(x)

2. Word-Level Attention

This layer computes the attention scores and generates word-level context vectors.

class WordAttention(nn.Module):
    def __init__(self, hidden_size):
        super(WordAttention, self).__init__()
        self.attention = nn.Sequential(
            nn.Linear(hidden_size, 1),
            nn.Softmax(dim=-1)
        )

    def forward(self, lstm_output):
        att_weights = self.attention(lstm_output)
        weighted_output = lstm_output * att_weights
        return weighted_output.sum(1), att_weights

3. Sentence Encoder & Attention

This is structurally similar to the word layers but operates at the sentence level.

Training the Model

With our data and model formats ready, we move on to training. Define a loss function like CrossEntropyLoss for classification tasks and an optimizer like Adam.

from torch.optim import Adam

model = HierarchicalAttentionNetwork(vocab_size, embed_size)
criterion = nn.CrossEntropyLoss()
optimizer = Adam(model.parameters(), lr=0.001)

Loop through epochs, passing the data through your model, calculate loss, and update the weights.

for epoch in range(num_epochs):
    for batch in train_iter:
        # Forward Pass
        outputs = model(batch.text)
        loss = criterion(outputs, batch.label)

        # Backward Pass
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()

    print(f'Epoch [{epoch+1}/{num_epochs}], Loss: {loss.item():.4f}')

Evaluation and Conclusion

Evaluate the model using test datasets to gauge its performance. Remember to calculate accuracy, precision, and recall to get a detailed understanding.

# Example evaluation snippet
def evaluate(model, test_loader):
    correct = 0
    total = 0
    with torch.no_grad():
        for data in test_loader:
            texts, labels = data
            outputs = model(texts)
            _, predicted = torch.max(outputs.data, 1)
            total += labels.size(0)
            correct += (predicted == labels).sum().item()
    return 100 * correct / total

With attention mechanisms taking the NLP field forward, utilizing hierarchical attention for document classification provides an intricate means to capture the nuances of language structure, significantly improving over traditional models.