memcnn 0.3.2

A PyTorch framework for developing memory efficient deep invertible networks.

A PyTorch framework for developing memory-efficient invertible neural networks.

• Free software: MIT license (please cite our work if you use it)

• Documentation: https://memcnn.readthedocs.io.

• Installation: https://memcnn.readthedocs.io/en/latest/installation.html

Features

• Simple ReversibleBlock wrapper class to wrap and convert arbitrary PyTorch Modules into invertible versions.

• Simple switching between additive and affine invertible coupling schemes and different implementations.

• Simple toggling of memory saving by setting the keep_input property of the ReversibleBlock.

• Training and evaluation code for reproducing RevNet experiments using MemCNN.

• CI tests for Python v2.7 and v3.6 and torch v0.4, v1.0, and v1.1 with good code coverage.

Example usage: ReversibleBlock

import torch
import torch.nn as nn
import memcnn


# define a new torch Module with a sequence of operations: Relu o BatchNorm2d o Conv2d
class ExampleOperation(nn.Module):
    def __init__(self, channels):
        super(ExampleOperation, self).__init__()
        self.seq = nn.Sequential(
                                    nn.Conv2d(in_channels=channels, out_channels=channels,
                                              kernel_size=(3, 3), padding=1),
                                    nn.BatchNorm2d(num_features=channels),
                                    nn.ReLU(inplace=True)
                                )

    def forward(self, x):
        return self.seq(x)

# generate some random input data (batch_size, num_channels, y_elements, x_elements)
X = torch.rand(2, 10, 8, 8)

# application of the operation(s) the normal way
model_normal = ExampleOperation(channels=10)
Y = model_normal(X)

# application of the operation(s) turned invertible using the reversible block
F = ExampleOperation(channels=10 // 2)
model_invertible = memcnn.ReversibleBlock(F, coupling='additive', keep_input=True, keep_input_inverse=True)
Y2 = model_invertible(X)

# The input (X) can be approximated (X2) by applying the inverse method of the reversible block on Y2
X2 = model_invertible.inverse(Y2)

Run PyTorch Experiments

After installing MemCNN run:

python -m memcnn.train [MODEL] [DATASET] [--fresh] [--no-cuda]

• Available values for DATASET are cifar10 and cifar100.

• Available values for MODEL are resnet32, resnet110, resnet164, revnet38, revnet110, revnet164

• Use the --fresh flag to remove earlier experiment results.

• Use the --no-cuda flag to train on the CPU rather than the GPU through CUDA.

Datasets are automatically downloaded if they are not available.

When using Python 3.* replace the python directive with the appropriate Python 3 directive. For example when using the MemCNN docker image use python3.6.