使用allegroai/clearml实现PyTorch文本分类任务全流程指南-优快云博客

本文链接：https://blog.youkuaiyun.com/gitblog_00099/article/details/148487999

使用allegroai/clearml实现PyTorch文本分类任务全流程指南

clearml ClearML - Auto-Magical CI/CD to streamline your ML workflow. Experiment Manager, MLOps and Data-Management 项目地址: https://gitcode.com/gh_mirrors/cl/clearml

项目概述

allegroai/clearml是一个开源的机器学习全生命周期管理平台，它能够帮助数据科学家和机器学习工程师更好地组织、跟踪、复制和协作机器学习实验。本文将详细介绍如何使用clearml平台实现一个基于PyTorch的文本分类任务，以AG News数据集为例，展示从环境配置到模型训练、评估和预测的完整流程。

环境准备与安装

在开始项目前，我们需要确保所有必要的依赖包已正确安装：

# 安装clearml核心包
pip install -U clearml>=0.15.0

# 安装PyTorch及相关依赖
pip install -U torch==1.5.0 torchtext==0.6.0 
pip install -U matplotlib==3.2.1 tensorboard==2.2.1

clearml的优势在于它能自动跟踪所有依赖包及其版本，确保实验的可复现性。安装完成后，我们导入必要的库：

import os
import time
import torch
import torch.nn as nn
from torchtext.datasets import text_classification
from torch.utils.tensorboard import SummaryWriter
from clearml import Task

初始化clearml任务

使用clearml的第一步是初始化一个任务，这将创建一个实验记录，自动跟踪所有代码、参数、模型和结果：

task = Task.init(project_name='Text Example', task_name='text classifier')
configuration_dict = {
    'number_of_epochs': 6, 
    'batch_size': 16, 
    'ngrams': 2, 
    'base_lr': 1.0
}
configuration_dict = task.connect(configuration_dict)

task.connect()方法特别重要，它允许我们在不修改代码的情况下通过clearml的Web界面动态调整超参数。这在超参数调优时非常有用。

数据准备与预处理

我们使用torchtext提供的AG News数据集，这是一个新闻文章分类数据集，包含4个类别：

if not os.path.isdir('./data'):
    os.mkdir('./data')
    
train_dataset, test_dataset = text_classification.DATASETS['AG_NEWS'](
    root='./data', 
    ngrams=configuration_dict.get('ngrams', 2)
vocabulary = train_dataset.get_vocab()

数据加载后，我们需要定义批处理函数：

def generate_batch(batch):
    label = torch.tensor([entry[0] for entry in batch])
    text = [entry[1] for entry in batch]
    offsets = [0] + [len(entry) for entry in text]
    offsets = torch.tensor(offsets[:-1]).cumsum(dim=0)
    text = torch.cat(text)
    return text, offsets, label

train_loader = torch.utils.data.DataLoader(
    train_dataset, 
    batch_size=configuration_dict.get('batch_size', 16),
    shuffle=True, 
    pin_memory=True, 
    collate_fn=generate_batch)

模型架构设计

我们构建一个简单的文本分类模型，包含嵌入层和全连接层：

class TextSentiment(nn.Module):
    def __init__(self, vocab_size, embed_dim, num_class):
        super().__init__()
        self.embedding = nn.EmbeddingBag(vocab_size, embed_dim, sparse=True)
        self.fc = nn.Linear(embed_dim, num_class)
        self.init_weights()

    def init_weights(self):
        initrange = 0.5
        self.embedding.weight.data.uniform_(-initrange, initrange)
        self.fc.weight.data.uniform_(-initrange, initrange)
        self.fc.bias.data.zero_()

    def forward(self, text, offsets):
        embedded = self.embedding(text, offsets)
        return self.fc(embedded)

模型初始化时，我们自动从数据集中获取词汇表大小和类别数：

VOCAB_SIZE = len(train_dataset.get_vocab())
EMBED_DIM = 32
NUM_CLASS = len(train_dataset.get_labels())
model = TextSentiment(VOCAB_SIZE, EMBED_DIM, NUM_CLASS).to(device)

训练与评估流程

训练过程通过clearml自动记录所有指标，并与TensorBoard集成：

criterion = torch.nn.CrossEntropyLoss().to(device)
optimizer = torch.optim.SGD(model.parameters(), lr=configuration_dict.get('base_lr', 1.0))
scheduler = torch.optim.lr_scheduler.StepLR(optimizer, 2, gamma=0.9)
tensorboard_writer = SummaryWriter('./tensorboard_logs')

def train_func(data, epoch):
    train_loss, train_acc = 0, 0
    for batch_idx, (text, offsets, cls) in enumerate(data):
        optimizer.zero_grad()
        output = model(text.to(device), offsets.to(device))
        loss = criterion(output, cls.to(device))
        loss.backward()
        optimizer.step()
        
        # 记录训练指标
        if batch_idx % 200 == 0:
            tensorboard_writer.add_scalar('training loss', loss, epoch*len(data)+batch_idx)
    
    scheduler.step()
    return train_loss / len(data.dataset), train_acc / len(data.dataset)

评估函数不仅计算指标，还会抽样展示模型预测结果：

def test(data, epoch):
    loss, acc = 0, 0
    for idx, (text, offsets, cls) in enumerate(data):
        with torch.no_grad():
            output = model(text.to(device), offsets.to(device))
            loss += criterion(output, cls.to(device)).item()
            acc += (output.argmax(1) == cls.to(device)).sum().item()
            
        # 记录样本预测结果
        if idx % 500 == 0:
            sample_text = ' '.join([vocabulary.itos[id] for id in text[offsets[0]:offsets[1]]])
            tensorboard_writer.add_text('Sample Prediction', sample_text, epoch)
    
    return loss / len(data.dataset), acc / len(data.dataset)

模型预测与部署

训练完成后，我们可以使用模型进行预测：

def predict(text, model, vocab, ngrams):
    tokenizer = get_tokenizer("basic_english")
    with torch.no_grad():
        text = torch.tensor([vocab[token] for token in ngrams_iterator(tokenizer(text), ngrams)])
        output = model(text, torch.tensor([0]))
        return output.argmax(1).item()

news_text = "MEMPHIS, Tenn. – Four days ago, Jon Rahm was..."
prediction = predict(news_text, model.to("cpu"), vocabulary, configuration_dict.get('ngrams', 2))
print(f"This is a {classes[prediction]} news")