Framework Migration Guide

https://chainer.github.io/migration-guide/

Authors: Chainer Team

This document provides technical information for migration from Chainer to PyTorch.

General Information

Concepts and components in both frameworks

Array Library

Chainer uses NumPy/CuPy (xp.ndarray) as an array library, and wraps them as chainer.Variable to support autograd. Similarly, PyTorch uses ATen (at::Tensor (C++)) as an array library ("tensor library" in PyTorch terms), and wraps it as torch::Tensor (C++ API) / torch.Tensor (Python API) to support autograd. torch.* provides API similar to (but not compatible with) NumPy, e.g. torch.dot, torch.float32, etc.

Core Framework and Training Loop

As both frameworks share the same concept, define-by-run, the look-and-feel of code written in PyTorch is pretty similar to Chainer. Here is the high-level mapping of features:

Refer to the Porting Guide section for the details of the difference of each component.

Migration scenarios

I want to port my Chainer script to PyTorch, step by step

Arguably the model is the hardest part to port without affecting the outcome of the training.

It might be easier to port in this order:

1. Training script (optimizer / updater / evaluator / ...)

• In order to use PyTorch optimizer to train a Chainer model, you will need cpm.LinkAsTorchModel.

2. Dataset / preprocessing

• Dataset is in general compatible between Chainer and PyTorch. This part can be delayed but also should be easy to do.

3. Model

• See the mapping of functions/modules below in this document.

I want to let my Chainer code train a PyTorch model

You can use cpm.TorchModule to wrap a PyTorch module as a Chainer model.

Migration tools (CPM)

chainer-pytorch-migration Python module (called CPM or "cpm" (module name) in this document) provides various utilities to help migration from Chainer to PyTorch.

Example code assumes that cpm is imported as follows:

cpm.TorchModule

This class wraps a PyTorch module as a Chainer link. It allows training PyTorch models in Chainer training scripts. The graph (forward/backward) must be constructed and traversed in PyTorch.

cpm.ChainerParameter

This class wraps a Chainer parameter as a PyTorch parameter. It allows training of Chainer models (chainer.Link) in PyTorch training scripts (with torch.optim.Optimizer). The graph (forward/backward) must be constructed and traversed in Chainer. cpm.LinkAsTorchModel internally uses it.

cpm.LinkAsTorchModel

This class automatically creates all the cpm.ChainerParameter objects for a given chainer link and provides methods such as parameters(), named_parameters() or state_dict() required by pytorch optimizers or tools such as horovod.

cpm.ignite.add_trainer_extension

This function registers a chainer trainer extension to be used with ignite.

Function call requires the ignite trainer, torch optimizer and the chainer extension as the parameters

pytorch-pfn-extras (ppe)

pytorch-pfn-extras Python module (called PPE or "ppe" (module name) in this document) provides various supplementary components for PyTorch, including APIs similar to Chainer, e.g. Extensions, Reporter, Lazy modules (automatically infer shapes of parameters). Here are some notable features Refer to the Documentation for the full list of features.

PPE also provides the interoperability feature between CuPy and PyTorch memory pool.

ppe.cuda.use_torch_mempool_in_cupy

This function makes CuPy use a memory pool from PyTorch. You need to call it before any operations using CuPy.

Note: The feature was originally implemented in CPM as cpm.use_torch_in_cupy_malloc, but has been moved to PPE. CPM version has been deprecated and not recommended any more.

Porting Guides

Dataset and data pre/post-processing

PyTorch datasets (pytorch.utils.data.Dataset) are basically compatible with Chainer’s. In most cases they are interchangeable in both directions.

Negative strides

As of PyTorch 1.2.0, PyTorch cannot handle data arrays with negative strides (can result from numpy.flip or chainercv.transforms.flip, for example).

Perhaps the easiest way to circumvent this problem is to wrap the dataset with numpy.ascontiguousarray.

Another way is to customize the collation function with collate_fn argument in torch.utils.data.DataLoader.

Rewriting Custom Converter Functions to collate_fn

In Chainer, it’s possible to specify custom converters for each batch via `training.updaters.StandardUpdater(train_iter, optimizer, device=device, converter=_converter)`. In PyTorch, similar functionality can be achieved via the data loader: `DataLoader(..., collate_fn=_converter)`.

There is, however, an important difference when used in conjunction with multiprocessing. In Chainer, `_converter` will be run in the main process, so it’s safe to access CUDA in the function when using multiprocessing’s `fork` mode. In PyTorch, however, `_converter` will be run inside each forker worker processes of the data loader. This means that we cannot access CUDA without getting a CUDA init error. It seems like in PyTorch, the correct usage is instead to only do CPU-related operations inside `_convert`, and only send the resulting tensors to the GPU *after* retrieving them from the data loader. The following is an example of correct PyTorch usage:

Note that the above scenario is different from what we expect in Chainer, where the `_converter` is called in the main process, which is why Chainer code might have CUDA-related operations inside the `_converter`.

Note that in the above use case, _convert should also use `pin_memory` in order to speed up the transfer of `(img, label)` from CPU to GPU: https://discuss.pytorch.org/t/when-to-set-pin-memory-to-true/19723

NumPy bridge

torch.DataLoader automatically converts NumPy arrays to PyTorch tensors, but if you want to do that manually, refer to NumPy Bridge.

CuPy bridge

DLPack can be used to bridge between CuPy and torch.Tensor. Note that DLPack does not handle ownership, so you have to make sure the original buffer (the original cupy.ndarray object or dltensor capsule object returned by toDlpack()) survives while the converted tensor/array is in use.

If you allocate a memory both in PyTorch and CuPy, it is also recommended to call ppe.cuda.use_torch_mempool_in_cupy before using CuPy to let CuPy use the PyTorch memory pool. Otherwise memories allocated and freed in CuPy will be kept in the CuPy memory pool which cannot be used by PyTorch.

Difference between PyTorch and NumPy/CuPy

Division Behavior

The behavior is different from NumPy/CuPy, which respects Python 3 division rules. You need to explicitly cast to float in PyTorch (discussion).

Feature Mapping

See PyTorch for Numpy users for the comparison table.

Training loop

This is an example code of training loop. Note model.train().

Evaluation loop

This is an example code of evaluation loop. Note model.eval() and with torch.no_grad().

Training and evaluation using Ignite

Ignite is something corresponding to chainer.training.Trainer in Chainer.

This Chainer code:

can be written in PyTorch using Ignite:

For a list of supported metrics, see https://pytorch.org/ignite/metrics.html.

Using Chainer extensions with Ignite

Using cpm.ignite.add_trainer_extension it is possible to register a chainer extension to be called within the ignite training loop.

A list of the supported extensions follows:

Some drawbacks rely on that metrics associated to the model or links might not accessible by default.

For example the user will need to report the loss or accuracy per iteration by using an ignite callback as this was done inside the chainer model.

Also for some extensions to work it is necessary for the user to assign the torch or chainer model to the optimizer target attribute and the output directory path for the LogReport, plotters and snapshot extensions

An example of how to register multiple extensions:

Snapshots

When using the snapshot extension with an ignite trainer, the pytorch objects are saved to an additional “snapshot-torch” file in the output folder. This allows to keep using these snapshots once the migration is finished and directly load pytorch models or the optimizer state from these files.

Additionally, if you are mixing chainer models or optimizers with ignite and pytorch, these objects will be saved in the chainer snapshot file.

The correct way to restore a snapshot is by using cpm.ignite.load_chainer_snapshot(engine, optimizer, snapshot_path) with the Chainer snapshot path.

Note that previously taken Chainer snapshots are not compatible.

Porting custom updater using Ignite

You can pass a step function to an Ignite engine.

• DCGAN example

Rewriting existing Chainer model

Use Mapping of functions and links to find and replace with the corresponding feature in PyTorch. You can also find existing model implementations in:

• PyTorch Examples
• torchvision.models module (CV models) (API Reference)
• PyTorch Hub (models from community) (API Reference)

Common pitfalls:

• Image format: RGB or BGR
• Image normalization: [0,1] or [0,255]
• Some old PyTorch examples and community projects are using torch.autograd.Variable, which is a deprecated interface.
• Some well-known models such as resnet might have different behavior in ChainerCV and torchvision. For example, chainercv.links.ResNet50 applies softmax to the output while torchvision.models.resnet50 does not.

Functions and Links

You can find the PyTorch equivalent of Chainer's functions and links in tables below.

Notes:

• Unlike NumPy/CuPy, PyTorch Tensor itself supports gradient computation (you can safely use torch.* or torch.nn.functional.* on torch.Tensor)
• Conventions of keyword arguments: dim and keepdim is used in PyTorch instead of axis and keepdims in Chainer/NumPy.
• Unlike Chainer, PyTorch provides the Module version of each function (e.g., nn.ReLU for F.relu), so you can use the Module version when defining model using functions in Sequential container.

Functions

F refers to chainer.functions (Chainer) / torch.nn.functional (PyTorch).

Links

L refers to chainer.links (Chainer), and nn refers to torch.nn (PyTorch).

Configuration

Here is the mapping of configurations in Chainer (chainer.config.*) and PyTorch:

See Reproducibility for the reproducibility (including steps to fix seeds).

Hooks

Function Hooks

There is no equivalent feature in PyTorch.

Replacements for Chainer built-in hooks:

• CUDAProfileHook: torch.autograd.profiler.profile + torch.autograd.profiler.emit_nvtx
• CupyMemoryProfileHook: N/A (allocator status can be retrieved)
• PrintHook: N/A
• TimerHook: torch.autograd.profiler.profile

Link Hooks

You can register Module Hooks per module. There's no way to inject a hook for every Module called under the specific scope.

Replacements for Chainer built-in hooks:

• SpectralNormalization: torch.nn.utils.spectral_norm / torch.nn.utils.remove_spectral_norm
• TimerHook: N/A

Optimizer Hooks

There is no direct equivalent in PyTorch, but you can register backward hooks per Tensor / Module to modify gradients.

Replacements for Chainer built-in hooks:

• WeightDecay: specify as weight_decay argument to each Optimizer (e.g., Adam)
• Lasso: N/A
• GradientClipping: torch.nn.utils.clip_grad_norm_
• GradientHardClipping: torch.nn.utils.clip_grad_value_
• GradientNoise: N/A
• GradientLARS: N/A (see pytorch #18414 and pytorch-lars)

Training PyTorch model using Chainer

To quickly try a PyTorch model in a training script using Chainer, cpm.TorchModule is the tool to use. Assuming you have a training script using Chainer, you have to try the following steps:

• Replace the model to train with cpm.TorchModule(module_you_want_to_use). Use to_gpu to transfer the variables to a GPU device.
• Rewrite the loss computation and backprop call with PyTorch.

• If you are using StandardUpdater, make its subclass and override update_core. Write loss calculation and backprop call in PyTorch.

• NOTE: Once you compute the gradient in PyTorch, it is automatically reflected to Chainer parameters, so it is valid to just call optimizer.update() after that.

• If you are using a custom training loop, rewrite the gradient computation part in PyTorch. See the above NOTE, too.

Distributed training

As of writing, there are two major ways to run distributed deep learning applications: torch.distributed and Horovod. We recommend torch.distributed as a first option because of the following reasons.

1. torch.distributed is a part of standard modules of PyTorch.
2. It supports some advanced features that Horovod doesn’t, such as multi-node batch normalization (e.g. inter-process batch normalization)

In this document, we describe both approaches to migrate ChainerMN programs to PyTorch.

Pytorch model using torch.distributed        

Torch.distributed is the standard module for distributed deep learning of PyTorch.

Torch.distributed supports three backends: “nccl”, “mpi” and “gloo”. For users who are migrating from Chainer and ChainerMN and have been using NCCL with MPI, using “nccl” backend is the most straightforward way. In this section, we assume that you use NCCL and MPI to run your distributed deep learning programs. In particular we assume Open MPI as the MPI implementation used here because it is the recommended option in ChainerMN, but other MPI implementations are mentioned as well.

Invocation

In ChainerMN, process invocation is totally coordinated by the MPI runtime. However, in PyTorch and torch.distributed, you may need a few more steps to invoke distributed deep learning processes. The simplest initialization method might be environment variable initialization.

The following environmental variables are necessary (whatever system you use to invoke your script, including MPI). Other variables, WORLD_SIZE and RANK, are set from inside the following snippet.

MASTER_ADDR : Address of the computing node where the rank 0 process runs.

MASTER_PORT : A free port of the MASTER_ADDR machine. The port will be used by  the rank 0 process.

Note that process invocation is highly system-dependent issue. PyTorch supports other options such as TCP initialization and shared file-system initialization. Please refer to the official documents for more details.

Initialization

The following code snippets shows how to initialize torch.distributed module.

Environmental variables set by MPI runtime are here, instead of communicator.intra_rank in ChainerMN because torch.distributed does not provide corresponding rank information. If you use MVAPICH2, use MV2_COMM_WORLD_SIZE, MV2_COMM_WORLD_RANK, MV2_COMM_WORLD_LOCAL_RANK respectively.

Dataset scattering

Each node can get a slice of a globally shared dataset using a DistributedSampler.

This will make every worker to only load a slice of the dataset, this sampler can be normally fed to the DataLoader.

Also, you need to call DistributedSampler.set_epoch() to adjust epoch numbers.  Thus typical training loop looks like:

Data transfer to devices

We need to specify the device to which the data is transferred using comm_local_rank.

Optimizer wrapping

In contrast to Horovod, We can use the same optimizer as in non-distributed execution.

Initial values broadcast

Parameter values are synchronized (i.e. initial broadcast and allreduce in every iteration) automatically by DistributedDataParallel class and thus no further modification is necessary.

Metrics average and reductions

https://pytorch.org/docs/stable/distributed.html#multi-gpu-collective-functions

Synchronization

To avoid potential data races other kinds of bugs, you may need to use torch.distributed.barrier() to synchronize processes before or after data loading, and finishing the application.

PyTorch model using Horovod

PyTorch can use Horovod to do Data Parallel training in a similar way to ChainerMN.
Data is distributed across the nodes and the optimizer is wrapped in with Horovod to automatically average the gradients of several MPI processes.

Horovod initialization

The following snippet shows how to import horovod and retrieve the current worker id and the total number of workers.

hvd.local_rank() is used to get the rank inside a single node, this is useful to assign GPUs, similar to ChainerMN’s intra_rank().

Dataset scattering

Each node can get a slice of a globally shared dataset using a DistributedSampler.

This will make every worker to only load a slice of the dataset, this sampler can be normally fed to the DataLoader

Optimizer wrapping

The optimizer is wrapped in a hvd.DistributedOptimizer object with the following configuration parameters.

compression : value in {hvd.Compression.fp16, hvd.Compression.none}

compression is used to reduce the size of the allreduce operations performed by the optimizer.

backward_passes_per_step : int default value usually 1

Number of batches that are performed locally before performing the gradients exchange.

From the documentation:

DistributedOptimizer exposes the synchronize() method, which forces allreduce operations to finish before continuing the execution. It’s useful in conjunction with gradient clipping, or other operations that modify gradients in place before step()is executed. Make sure to use optimizer.skip_synchronize() if you’re calling synchronize() in your code.

Initial values broadcast

Before starting the training loop, initial model parameters and the optimizer state must be broadcasted to all the workers:

Metrics average and reductions

When computing the loss and other metrics such as accuracy, the values of multiple workers can be explicitly exchanged to compute averages:

Horovod has support to exchange data using other MPI collectives:

• horovod.torch.allgather
• horovod.torch.broadcast

There are _async versions of the three functions that can be queried using poll() on the returned handler or synchronize() to wait till completion.

Horovod code structure

Obtaining Horovod traces to measure performance

Communication traces showing Horovod communications can be obtained by setting the HOROVOD_TIMELINE environment variable.

The resultant trace can be visualized in Chrome by using the browser built-in chrome://tracing feature.

Tuning Horovod performance

Horovod has several knobs to improve its performance

• TensorFusion improves network utilization when there are many operations to be done on small tensors. Instead of eagerly starting multiple small reductions, horovod combines the small tensors on a buffer and then batch operations in order to perform a more efficient communication/computation overlap. HOROVOD_FUSION_THRESHOLD controls the size of the buffer in bytes (default is 64MB). HOROVOD_CYCLE_TIME specifies the amount of time in ms to wait for tensors to be batched before reducing them.
• Hierarchical Reduce: HOROVOD_HIERARCHICAL_ALLREDUCE does a hybrid approach where allreduce in-node is done by NCCL, and allreduce across nodes is done by MPI.
• HOROVOD_AUTOTUNE can be used for horovod to perform a Bayesian optimization on all the configurable parameters at the first epochs of training. More info here

Using Horovod with apex

Horovod launches all-reduce in parallel with backward computation, and apex unscales gradient after backward computation.

To avoid race conditions, we have to wait for all-reduce completion before unscaling:

Also, backward_passes_per_step should be 1 when using Horovod and apex. The current implementation of Horovoda and apex do not work as expected when backward_passes_per_step is not 1.

Multi-Node Batch Normalization in Horovod

Horovod does not yet officially support MNBN (https://github.com/horovod/horovod/issues/1384), but there exists an unofficial implementation: https://github.com/atranitell/Synchronized-BatchNorm-PyTorch-Horovod/blob/master/sync_bn.py. Apex also has an implementation: https://nvidia.github.io/apex/parallel.html#apex.parallel.SyncBatchNorm

Gathering arbitrary objects using Horovod and mpi4py

Horovod supports simultaneous usage with mpi4py (https://github.com/horovod/horovod#mpi4py). You can directly work with mpi4py to e.g. rewrite ChainerMN's comm.gather_obj:

Alternatives to Horovod

Horovod is introduced here because it greatly resembles ChainerMN and can be used in our computing infrastructure right away. Alternatives are:

• Pytorch distributed API
• Lightning

Chainer model using Horovod

To train chainer models in distributed environments using Horovod, the chainer link should be wrapped using cpm.LinkAsTorchModel. The use of a PyTorch optimizer is required.

PyTorch model using ChainerMN

Using the cpm tool it is also possible to train a PyTorch model using ChainerMN.

The current support is limited only to data parallel training.

Porting code that edits the computational graph

Unchaining nodes

Explains differences of how variables can be unchained from the computational graph.

• Chainer: Variable.{data,array} or Variable.unchain().
• PyTorch: Tensor.detach(). This method returns a new Tensor unchained from the computational graph with requires_grad set to False. This is not an in-place operation in contrast to Variable.unchain() and does not incur any side effects (although the new Tensor will share the same memory). An in-place variant is however provided as well Tensor.detach_().
• PyTorch: Tensor.data is discouraged and it seems like it might even get deprecated in the future (based on comments in forums and on GitHub).
• Chainer’s Variable.unchain_backward() counterpart in PyTorch is not available?

Backprop modes

Explains differences of how backprop modes are switched.

• Chainer: no_backprop_mode, force_backprop_mode correspond to the following in PyTorch.
• PyTorch: torch.autograd.no_grad(), torch.autograd.enable_grad().
• no_grad() can also be used to allow writing data directly to an already allocated Tensor using [...].
• In Chainer, this is also a configuration (configuration.config.enable_backprop), it is however not in PyTorch.
• Both Chainer and PyTorch default to backprop mode being enabled.

Train/Test modes

Explains differences of how train/test modes are switched.

• Chainer: This mode is controlled via a configuration (configuration.config.train).
• PyTorch: This mode is a module state and should be changed using torch.nn.Module.train(mode), torch.nn.Module.eval(). This also means that you must pass the “train” argument to functions calls such as torch.nn.functional.dropout otherwise use the torch.nn.Dropout module which is stateful.
• In both Chainer and PyTorch, train/test mode affects the behavior of certain functions/links/modules such as dropout and batch normalization.
• Both Chainer and PyTorch default to train mode.

Ecosystem

This section introduces some of the larger repositories under the PyTorch GitHub organization. It also refers to the official list of other ecosystem-libraries acknowledged by PyTorch.

PyTorch

Summary: Tensors and Dynamic neural networks in Python with strong GPU acceleration

GitHub: https://github.com/pytorch/pytorch

• Installation: See Start Locally for the instructions. Note that not all variants are hosted on PyPI; as of PyTorch 1.2.0, torch/torchvision packages hosted on PyPI are for CUDA 10.0.
• Nightly builds: Nightly build is provided so you can try the pre-built binary with new merged feature next day. Each nightly build is tagged per each day so you can download nightly version locked pip wheel and prebuilt libtorch of specific build date. For pip wheel refer Start Locally (CPU page, CUDA 9.2 page, CUDA 10.0 page) to get optimal version.

Ignite

Summary: High level utilities such as training loop abstraction.

GitHub: https://github.com/pytorch/ignite

torchvision

Summary: PyTorch for CV.

GitHub: https://github.com/pytorch/vision

Recommended by the official installation guide to install along with pytorch.

Provides domain-agnostic (not limited to CV) data augmentation functionality.

Provides loaders for video data. Slow due to ffmpeg but this might be improved in the future?

torchtext

Summary: PyTorch for NLP.

GitHub: https://github.com/pytorch/text

torchaudio

Summary: PyTorch for audio data.

GitHub: https://github.com/pytorch/audio

Fairseq

Summary: Seq2seq models.

GitHub: https://github.com/pytorch/fairseq

Seq2seq models such as translation. Includes the Transformer and BERT-like models.

Other

There is an official list of libraries included in the PyTorch ecosystem (besides the domain specific libraries above), including e.g. Ignite.

https://pytorch.org/ecosystem