torchaid 0.1.4

A PyTorch utility library that streamlines the deep learning training pipeline.

torchaid

torchaid is a PyTorch utility library that provides structured abstractions and reusable components to streamline the deep learning training pipeline.

Features

• Structured training abstractions — type-safe base classes for inputs, outputs, metrics, and settings built on Pydantic v2
• Training framework — a full training loop with mixed-precision support, automatic checkpointing, metric logging (CSV), and early stopping
• Transformer modules — standard and relative-position-aware Transformer encoder layers with multi-head self-attention
• Task templates — ready-to-use implementation for multi-label classification
• Utilities — dataset splitting, random seed management, attention mask generation, and JSON-to-Pydantic loading
• Learning rate schedulers — cosine decay with linear warm-up, and triangular2 cyclic scheduling

Requirements

• Python 3.10+
• PyTorch 2.0+

Installation

pip install torchaid

Or install from source:

git clone https://github.com/harunori-kawano/torchaid.git
cd torchaid
pip install -e .

Quick Start

1. Define your data schema

from torchaid import BaseInputs, BaseOutputs
import torch

class MyInputs(BaseInputs):
    input_ids: torch.Tensor   # (B, L)
    labels: torch.LongTensor  # (B,)

class MyOutputs(BaseOutputs):
    logits: torch.Tensor      # (B, num_classes)

2. Implement your model

from torchaid import TaskModule, Mode, BaseOutputs
from torch import nn

class MyModel(TaskModule):
    def __init__(self, vocab_size: int, num_classes: int):
        super().__init__()
        self.embed = nn.Embedding(vocab_size, 128)
        self.classifier = nn.Linear(128, num_classes)
        self.criterion = nn.CrossEntropyLoss()

    def forward(self, mode: Mode, batch: MyInputs) -> BaseOutputs:
        x = self.embed(batch.input_ids).mean(dim=1)
        logits = self.classifier(x)
        loss = self.criterion(logits, batch.labels)
        return MyOutputs(loss=loss, logits=logits)

3. Define metrics and settings

from torchaid import BaseMetrics, BaseSettings, BaseMetricCalculator
from typing import Optional, Any

class MyMetrics(BaseMetrics):
    train_loss: Optional[float] = None
    val_loss: Optional[float] = None

class MySettings(BaseSettings):
    batch_size: int = 32
    max_epoch_num: int = 10

class MyCalculator(BaseMetricCalculator[MyMetrics]):
    def __init__(self):
        super().__init__(MyMetrics())
        self._losses: list[float] = []

    def train_step(self, outputs, batch) -> dict[str, Any]:
        loss = outputs.loss.item()
        self._losses.append(loss)
        return {"loss": loss}

    def val_step(self, outputs, batch) -> dict[str, Any]:
        return self.train_step(outputs, batch)

    def test_step(self, outputs, batch) -> dict[str, Any]:
        return self.train_step(outputs, batch)

    def check(self) -> bool:
        import statistics
        self.metrics.train_loss = statistics.mean(self._losses)
        return True

    def test(self): pass

    def reset(self):
        self._losses.clear()

4. Train

import torch
from torchaid.core.trainer import TrainFramework

settings = MySettings(batch_size=32, max_epoch_num=10, device="cuda")
model = MyModel(vocab_size=1000, num_classes=5)
optimizer = torch.optim.AdamW(model.parameters(), lr=1e-4)

framework = TrainFramework(
    model=model,
    ls=settings,
    metric_calculator=MyCalculator(),
    optimizer=optimizer,
    inputs_config=MyInputs,
)

framework.train(train_dataset, val_dataset, save_dir="./outputs")

Module Overview

Module	Description
torchaid.core	Base classes (BaseInputs, BaseOutputs, BaseMetrics, BaseSettings, TaskModule, BaseMetricCalculator, Mode) and TrainFramework
torchaid.templates.multilabel_classification	Complete template for multi-label classification
torchaid.extras.modules.transformer	Transformer, TransformerWithRelativePosition, and sub-modules
torchaid.extras.modules.positional_encoders	PositionalEmbedding, RelativePositionEmbedding
torchaid.extras.utils	split_dataset, set_random_seed, make_attention_mask, json_to_instance
torchaid.extras.scheduler	get_cosine_scheduler, get_cycle_scheduler

Template: Multi-Label Classification

from torchaid.templates import multilabel_classification as mlc
from torchaid.core.trainer import TrainFramework
from torch import nn
import torch

backbone = nn.Sequential(nn.Linear(128, 64), nn.ReLU(), nn.Linear(64, 10))
model = mlc.MultiLabelClassification(backbone)
optimizer = torch.optim.Adam(model.parameters())

framework = TrainFramework(
    model=model,
    ls=settings,
    metric_calculator=mlc.MetricsCalculator(),
    optimizer=optimizer,
    inputs_config=mlc.Inputs,
)

Extras

Cosine Decay Scheduler

from torchaid.extras.scheduler import get_cosine_scheduler

scheduler = get_cosine_scheduler(
    optimizer, warmup_steps=500, max_steps=10000
)

Dataset Split

from torchaid.extras.utils import split_dataset

train, val, test = split_dataset(dataset, ratios=[8, 1, 1], seed=42)

Relative Position Transformer

from torchaid.extras.modules.transformer import TransformerWithRelativePosition

layer = TransformerWithRelativePosition(
    hidden_size=256,
    intermediate_size=1024,
    num_attention_heads=8,
    dropout_probability=0.1,
    max_length=512,
    with_cls=True,
)

License

MIT License. See LICENSE for details.