A PyTorch framework for developing memory efficient deep invertible networks.
Project description
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.
Project details
Release history Release notifications | RSS feed
Download files
Download the file for your platform. If you're not sure which to choose, learn more about installing packages.
Source Distribution
Built Distribution
Hashes for memcnn-0.3.3-py2.py3-none-any.whl
Algorithm | Hash digest | |
---|---|---|
SHA256 | f542c48e7f9054982d28b404878198003d85a98b1437a96b176ec7ba65fbde47 |
|
MD5 | d766ac91f7f1eea59ad20a13dfe9b312 |
|
BLAKE2b-256 | 9517f5ce7d0f827604c20810411408a39018a29c58c82211371997f5b1d170b4 |