Skip to main content

Framework for creating (partially) reversible neural networks with PyTorch

Project description


Framework for creating (partially) reversible neural networks with PyTorch

RevTorch is introduced and explained in our paper A Partially Reversible U-Net for Memory-Efficient Volumetric Image Segmentation, which was accepted for presentation at MICCAI 2019.

If you find this code helpful in your research please cite the following paper:

         author={Br{\"u}gger, Robin and Baumgartner, Christian F.
         and Konukoglu, Ender},
         title={A Partially Reversible U-Net for Memory-Efficient Volumetric Image Segmentation},


Use pip to install RevTorch:

$ pip install revtorch

RevTorch requires PyTorch. However, PyTorch is not included in the dependencies since the required PyTorch version is dependent on your system. Please install PyTorch following the instructions on the PyTorch website.


This example shows how to use the RevTorch framework.

import torch
import torchvision
import torchvision.transforms as transforms
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim
import revtorch as rv

def train():
    trainset = torchvision.datasets.CIFAR10(root="./data", train=True, download=True, transform=transforms.ToTensor())
    trainloader =, batch_size=4, shuffle=True)

    net = PartiallyReversibleNet()
    criterion = nn.CrossEntropyLoss()
    optimizer = optim.Adam(net.parameters())

    for epoch in range(2):

        running_loss = 0.0
        for i, data in enumerate(trainloader, 0):
            inputs, labels = data
            outputs = net(inputs)
            loss = criterion(outputs, labels)

            #logging stuff
            running_loss += loss.item()
            LOG_INTERVAL = 200
            if i % LOG_INTERVAL == (LOG_INTERVAL-1):  # print every 2000 mini-batches
                print('[%d, %5d] loss: %.3f' % (epoch + 1, i + 1, running_loss / LOG_INTERVAL))
                running_loss = 0.0

class PartiallyReversibleNet(nn.Module):
    def __init__(self):
        super(PartiallyReversibleNet, self).__init__()

        #initial non-reversible convolution to get to 32 channels
        self.conv1 = nn.Conv2d(3, 32, 3)

        #construct reversible sequencce with 4 reversible blocks
        blocks = []
        for i in range(4):

            #f and g must both be a nn.Module whos output has the same shape as its input
            f_func = nn.Sequential(nn.ReLU(), nn.Conv2d(16, 16, 3, padding=1))
            g_func = nn.Sequential(nn.ReLU(), nn.Conv2d(16, 16, 3, padding=1))

            #we construct a reversible block with our F and G functions
            blocks.append(rv.ReversibleBlock(f_func, g_func))

        #pack all reversible blocks into a reversible sequence
        self.sequence = rv.ReversibleSequence(nn.ModuleList(blocks))

        #non-reversible convolution to get to 10 channels (one for each label)
        self.conv2 = nn.Conv2d(32, 10, 3)

    def forward(self, x):
        x = self.conv1(x)

        #the reversible sequence can be used like any other nn.Module. Memory-saving backpropagation is used automatically
        x = self.sequence(x)

        x = self.conv2(F.relu(x))
        x = F.avg_pool2d(x, (x.shape[2], x.shape[3]))
        x = x.view(x.shape[0], x.shape[1])
        return x

if __name__ == "__main__":

Python version

Tested with Python 3.6 and PyTorch 1.1.0. Should work with any version of Python 3.


Version 0.2.4

  • Added option to disable eager discarding of variables to allow for multiple backward() calls

Version 0.2.3

  • Added option to use the same random seed for the forward and backwar pass (Pull request)

Version 0.2.1

  • Added option to select the dimension along which the tensor is split (Pull request)

Version 0.2.0

  • Fixed memory leak when not consuming output of the reversible block (Issue)

Version 0.1.0

  • Initial release

Project details

Download files

Download the file for your platform. If you're not sure which to choose, learn more about installing packages.

Source Distribution

revtorch-0.2.4.tar.gz (5.6 kB view hashes)

Uploaded source

Supported by

AWS AWS Cloud computing Datadog Datadog Monitoring Facebook / Instagram Facebook / Instagram PSF Sponsor Fastly Fastly CDN Google Google Object Storage and Download Analytics Huawei Huawei PSF Sponsor Microsoft Microsoft PSF Sponsor NVIDIA NVIDIA PSF Sponsor Pingdom Pingdom Monitoring Salesforce Salesforce PSF Sponsor Sentry Sentry Error logging StatusPage StatusPage Status page