anandbhattad / inferno

A utility library around PyTorch

Inferno

Inferno is a little library providing utilities and convenience functions/classes around PyTorch. It's a work-in-progress, but the latest release (0.3.1) should be fairly stable!

• Free software: Apache Software License 2.0
• Documentation: http://inferno-pytorch.readthedocs.io (Work in Progress).

Features

Current features include:

• a basic Trainer class to encapsulate the training boilerplate (iteration/epoch loops, validation and checkpoint creation),
• a graph API for building models with complex architectures, powered by networkx.
• easy data-parallelism over multiple GPUs,
• a submodule for torch.nn.Module-level parameter initialization,
• a submodule for data preprocessing / transforms,
• support for Tensorboard (best with atleast tensorflow-cpu installed)
• a callback API to enable flexible interaction with the trainer,
• various utility layers with more underway,
• a submodule for volumetric datasets, and more!

import torch.nn as nn
from inferno.io.box.cifar import get_cifar10_loaders
from inferno.trainers.basic import Trainer
from inferno.trainers.callbacks.logging.tensorboard import TensorboardLogger
from inferno.extensions.layers.convolutional import ConvELU2D
from inferno.extensions.layers.reshape import Flatten

# Fill these in:
LOG_DIRECTORY = '...'
SAVE_DIRECTORY = '...'
DATASET_DIRECTORY = '...'
DOWNLOAD_CIFAR = True
USE_CUDA = True

# Build torch model
model = nn.Sequential(
    ConvELU2D(in_channels=3, out_channels=256, kernel_size=3),
    nn.MaxPool2d(kernel_size=2, stride=2),
    ConvELU2D(in_channels=256, out_channels=256, kernel_size=3),
    nn.MaxPool2d(kernel_size=2, stride=2),
    ConvELU2D(in_channels=256, out_channels=256, kernel_size=3),
    nn.MaxPool2d(kernel_size=2, stride=2),
    Flatten(),
    nn.Linear(in_features=(256 * 4 * 4), out_features=10),
    nn.LogSoftmax(dim=1)
)

# Load loaders
train_loader, validate_loader = get_cifar10_loaders(DATASET_DIRECTORY,
                                                    download=DOWNLOAD_CIFAR)

# Build trainer
trainer = Trainer(model) \
  .build_criterion('NLLLoss') \
  .build_metric('CategoricalError') \
  .build_optimizer('Adam') \
  .validate_every((2, 'epochs')) \
  .save_every((5, 'epochs')) \
  .save_to_directory(SAVE_DIRECTORY) \
  .set_max_num_epochs(10) \
  .build_logger(TensorboardLogger(log_scalars_every=(1, 'iteration'),
                                  log_images_every='never'),
                log_directory=LOG_DIRECTORY)

# Bind loaders
trainer \
    .bind_loader('train', train_loader) \
    .bind_loader('validate', validate_loader)

if USE_CUDA:
  trainer.cuda()

# Go!
trainer.fit()

To visualize the training progress, navigate to LOG_DIRECTORY and fire up tensorboard with

$ tensorboard --logdir=${PWD} --port=6007

and navigate to localhost:6007 with your browser.

Installation

Conda packages for python >= 3.6 for all distributions are availaible on conda-forge:

$ conda install -c pytorch -c conda-forge inferno

Future Features:

Planned features include:

• a class to encapsulate Hogwild! training over multiple GPUs,
• minimal shape inference with a dry-run,
• proper packaging and documentation,
• cutting-edge fresh-off-the-press implementations of what the future has in store. :)

Credits

All contributors are listed here_. .. _here: https://inferno-pytorch.github.io/inferno/html/authors.html

This package was partially generated with Cookiecutter and the audreyr/cookiecutter-pypackage project template + lots of work by Thorsten.