Skip to main content

Torchium

PyPI version Python License: MIT Downloads Build Status codecov Documentation Status Code style: black Imports: isort

Torchium is the most comprehensive PyTorch extension library, providing 65+ advanced optimizers and 70+ specialized loss functions for deep learning research and production. Built on top of PyTorch's robust foundation, Torchium seamlessly integrates cutting-edge optimization algorithms and loss functions from various domains including computer vision, natural language processing, generative models, and metric learning.

Key Features

  • Advanced Optimizers: 65+ state-of-the-art optimizers including Lion, Ranger21, AdaBelief, SAM, and experimental evolutionary algorithms
  • Specialized Loss Functions: 70+ domain-specific loss functions including Focal Loss, Dice Loss, Perceptual Loss, and advanced metric learning losses
  • Domain Expertise: Specialized optimizers and losses for computer vision, NLP, generative models, medical imaging, and audio processing
  • Performance Optimized: Highly optimized implementations with full CUDA support and memory efficiency
  • Seamless Integration: Drop-in replacement for PyTorch optimizers and losses with zero breaking changes
  • Research-Grade: Latest optimization techniques from top-tier machine learning papers
  • Production Ready: Comprehensive documentation, extensive testing, and robust error handling

Quick Start

Installation

pip install torchium

Basic Usage

import torch
import torch.nn as nn
import torchium

# Create a model
model = nn.Sequential(
    nn.Linear(10, 64),
    nn.ReLU(),
    nn.Linear(64, 1)
)

# Use any of the 65+ optimizers
optimizer = torchium.optimizers.Ranger(model.parameters(), lr=1e-3)
# optimizer = torchium.optimizers.Lion(model.parameters(), lr=1e-4)
# optimizer = torchium.optimizers.AdaBelief(model.parameters(), lr=1e-3)

# Use any of the 70+ loss functions
criterion = torchium.losses.FocalLoss(alpha=0.25, gamma=2.0)
# criterion = torchium.losses.DiceLoss(smooth=1e-5)
# criterion = torchium.losses.PerceptualLoss()

# Training loop
for data, target in dataloader:
    optimizer.zero_grad()
    output = model(data)
    loss = criterion(output, target)
    loss.backward()
    optimizer.step()

Factory Functions

import torchium

# Create optimizers using factory functions
optimizer = torchium.create_optimizer('ranger', model.parameters(), lr=1e-3)
criterion = torchium.create_loss('focal', alpha=0.25, gamma=2.0)

# List available optimizers and losses
print("Available optimizers:", torchium.get_available_optimizers())
print("Available losses:", torchium.get_available_losses())

Optimizer Categories

Adaptive Optimizers

  • Adam Family: Adam, AdamW, RAdam, AdaBelief, AdaBound, AdaHessian, AdamP, AdamS, AdamD
  • Adagrad Family: Adagrad, Adadelta, AdaFactor, AdaGC, AdaGO, AdaLOMO, Adai, Adalite
  • RMSprop Family: RMSprop, Yogi

Momentum-Based Optimizers

  • SGD Variants: SGD, NesterovSGD, QHM, AggMo, SWATS, SGDP, SGDSaI, SignSGD
  • Classical: HeavyBall, NAG (Nesterov Accelerated Gradient)

Specialized Optimizers

  • Computer Vision: Ranger, Ranger21, Ranger25, AdamP
  • NLP: LAMB, NovoGrad, AdaFactor
  • Large Scale: LARS, LAMB (Layer-wise Adaptive Moments optimizer)
  • Memory Efficient: Lion, MADGRAD, SM3

Meta-Optimizers

  • Sharpness-Aware: SAM, GSAM, ASAM, LookSAM, WSAM
  • Gradient Methods: Lookahead, GradientCentralization, PCGrad

Second-Order Optimizers

  • Quasi-Newton: LBFGS, Shampoo, AdaHessian
  • Natural Gradients: K-FAC, Natural Gradient Descent

Experimental Optimizers

  • Evolutionary: CMA-ES, Differential Evolution, Particle Swarm Optimization
  • Quantum-Inspired: Quantum Annealing
  • Genetic: Genetic Algorithm

Loss Function Categories

Classification Losses

  • Cross-Entropy Variants: Standard, Focal, Label Smoothing, Class-Balanced
  • Margin-Based: Triplet, Contrastive, Angular, ArcFace, CosFace
  • Ranking: NDCG, MRR, MAP, RankNet, LambdaRank

Computer Vision Losses

  • Segmentation: Dice, IoU, Tversky, Focal Tversky, Lovász, Boundary
  • Object Detection: Focal, GIoU, DIoU, CIoU, EIoU, α-IoU
  • Super-Resolution: Perceptual, SSIM, MS-SSIM, LPIPS, VGG
  • Style Transfer: Style, Content, Total Variation

NLP Losses

  • Language Modeling: Perplexity, Cross-entropy variants
  • Sequence Labeling: CRF, Structured prediction
  • Text Generation: BLEU, ROUGE, METEOR, BERTScore

Generative Model Losses

  • GAN: Standard, Wasserstein, Hinge, Least Squares, Relativistic
  • VAE: ELBO, β-VAE, β-TC-VAE, Factor-VAE
  • Diffusion: DDPM, DDIM, Score matching

Regression Losses

  • Standard: MSE, MAE, Huber, Quantile, Log-cosh
  • Robust: Tukey, Cauchy, Welsch, Fair

Metric Learning Losses

  • Contrastive: Contrastive, Triplet, Quadruplet, N-Pair
  • Angular: Angular, ArcFace, CosFace, SphereFace
  • Proxy: ProxyNCA, ProxyAnchor

Multi-task Learning

  • Uncertainty Weighting: Automatic task balancing
  • Gradient Balancing: PCGrad, GradNorm, CAGrad
  • Dynamic Balancing: Adaptive task weighting

Domain-Specific Examples

Computer Vision

import torchium

# Segmentation with Dice + Focal Loss
criterion = torchium.losses.CombinedSegmentationLoss(
    dice_weight=0.5, 
    focal_weight=0.5
)

# Object Detection with GIoU Loss
bbox_loss = torchium.losses.GIoULoss()

# Super-Resolution with Perceptual Loss
perceptual_loss = torchium.losses.PerceptualLoss(
    feature_layers=['conv_4_2', 'conv_5_2']
)

Natural Language Processing

# BERT Training with LAMB optimizer
optimizer = torchium.optimizers.LAMB(
    model.parameters(), 
    lr=1e-3, 
    weight_decay=0.01
)

# Sequence-to-Sequence with Label Smoothing
criterion = torchium.losses.LabelSmoothingLoss(smoothing=0.1)

Generative Models

# GAN Training
g_optimizer = torchium.optimizers.Ranger21(generator.parameters(), lr=2e-4)
d_optimizer = torchium.optimizers.Ranger21(discriminator.parameters(), lr=2e-4)

# Wasserstein GAN Loss
gan_loss = torchium.losses.WassersteinLoss()

# VAE Training with β-VAE Loss
vae_loss = torchium.losses.BetaVAELoss(beta=4.0)

Medical Imaging

# Medical image segmentation
medical_loss = torchium.losses.MedicalImagingLoss(
    dice_weight=0.6,
    ce_weight=0.4
)

# Optimizer for medical imaging
optimizer = torchium.optimizers.AdaBelief(
    model.parameters(),
    lr=1e-3,
    weight_decay=1e-4
)

Audio Processing

# Audio enhancement with multi-scale loss
audio_loss = torchium.losses.AudioProcessingLoss(
    time_weight=0.7,
    freq_weight=0.3
)

# Optimizer for audio tasks
optimizer = torchium.optimizers.NovoGrad(
    model.parameters(),
    lr=1e-3
)

Performance Benchmarks

Torchium optimizers are benchmarked against PyTorch's built-in optimizers:

Optimizer CIFAR-10 Accuracy Training Time Memory Usage
SGD 91.2% 100% 100%
Adam 92.8% 98% 105%
Ranger21 94.1% 95% 102%
AdaBelief 93.6% 97% 103%
Lion 93.4% 88% 85%

Advanced Features

Meta-Optimization with SAM

import torchium

# Sharpness-Aware Minimization
base_optimizer = torch.optim.SGD
sam_optimizer = torchium.optimizers.SAM(
    model.parameters(),
    base_optimizer,
    rho=0.05
)

# Training with SAM
for data, target in dataloader:
    def closure():
        sam_optimizer.zero_grad()
        output = model(data)
        loss = criterion(output, target)
        loss.backward()
        return loss
    
    sam_optimizer.step(closure)

Multi-task Learning

# Uncertainty weighting for multi-task learning
mtl_loss = torchium.losses.UncertaintyWeightingLoss(num_tasks=3)

# PCGrad for conflicting gradients
pcgrad_optimizer = torchium.optimizers.PCGrad(
    model.parameters(),
    torch.optim.Adam,
    num_tasks=3
)

Experimental Optimizers

# Evolutionary optimization
cmaes_optimizer = torchium.optimizers.CMAES(
    model.parameters(),
    sigma=0.1,
    popsize=50
)

# Particle Swarm Optimization
pso_optimizer = torchium.optimizers.ParticleSwarmOptimization(
    model.parameters(),
    popsize=30,
    inertia=0.9
)

Documentation [To be updated]

Contributing

We welcome contributions! Please see our Contributing Guide for details.

Adding New Optimizers/Losses

import torchium
from torchium.utils.registry import register_optimizer, register_loss

@register_optimizer("my_optimizer")
class MyOptimizer(torch.optim.Optimizer):
    def __init__(self, params, lr=1e-3):
        defaults = dict(lr=lr)
        super().__init__(params, defaults)
    
    def step(self, closure=None):
        # Implementation here
        pass

@register_loss("my_loss")
class MyLoss(torch.nn.Module):
    def __init__(self):
        super().__init__()
    
    def forward(self, input, target):
        # Implementation here
        pass

License

This project is licensed under the MIT License - see the LICENSE file for details.

Citation

If you use Torchium in your research, please cite:

@software{torchium2025,
    title={Torchium: Advanced PyTorch Extension Library},
    author={Vishesh Yadav},
    year={2025},
    url={https://github.com/vishesh9131/torchium},
    version={0.1.0}
}

Acknowledgments

  • Built on top of PyTorch
  • Inspired by pytorch-optimizer
  • Thanks to all the researchers who developed these optimization algorithms

Support


Made with dedication by the @vishesh9131

Supercharge your PyTorch models with the most comprehensive collection of optimizers and loss functions!

Metadata

Release files for torchium 0.1.5

For a detailed explanation of source distributions (sdists) and built distributions (wheels), please see the package formats documentation.

Source distribution (sdist)

Source distribution for torchium 0.1.5
File Size Uploaded Provenance
torchium-0.1.5.tar.gz 2.2 MB Without provenance Details

Built distribution (wheel)

Table of built distributions (wheels) for torchium 0.1.5
File Interpreter ABI Platform
torchium-0.1.5-py3-none-any.whl Python 3 none any Details

Total release size: 2.7 MB

Release files / torchium-0.1.5.tar.gz

Download URL torchium-0.1.5.tar.gz
Size 2.2 MB
Tags Source
SHA-256 checksum
How to use checksums
61c6a206b1fba76932cd4a353fa6ff986bc4c61a03213ee994ecb689eb9348b9
BLAKE2b-256 checksum
How to use checksums
889ffa39eaa06f5f290708fa7b414861f2e46eb132b2b5cc8f00a87b8731c818
Upload date
Uploaded using Trusted Publishing?
What is trusted publishing?
No
Uploaded via twine/6.2.0 CPython/3.11.4

Release files / torchium-0.1.5-py3-none-any.whl

Download URL torchium-0.1.5-py3-none-any.whl
Size 486.9 kB
Tags Python 3
SHA-256 checksum
How to use checksums
97df5001605da381424411d35071c39cfa5cb51db0c2af7ce4270cc2ad448ff3
BLAKE2b-256 checksum
How to use checksums
db78360ed968d333f63127724c0d329b5bc68616e1865e30fa160b5810a4d3dd
Upload date
Uploaded using Trusted Publishing?
What is trusted publishing?
No
Uploaded via twine/6.2.0 CPython/3.11.4

Security for torchium 0.1.5

Provenance

Provenance describes where a file came from. On PyPI, provenance is shared via attestations, which provide a verifiable record of the build or publishing details. View details, limitations and caveats.

No provenance found

None of the files in this release were uploaded with provenance metadata.

A note on package security

PyPI is a repository for distributing packages, not a service for auditing them. While third-party security researchers and our staff do their best to detect and remove malicious packages, PyPI cannot guarantee the quality, security or suitability of any given package.

To learn more about protecting yourself from malicious actors, visit our package security documentation.

Anthropic, PBC Visionary sponsor Bloomberg Visionary sponsor Hudson River Trading Visionary sponsor Meta Visionary sponsor NVIDIA Visionary sponsor Microsoft Sustainability sponsor Depot Continuous Integration AWS Cloud computing and Security Sponsor Datadog Monitoring Fastly CDN Google Download Analytics Sentry Error logging StatusPage Status page