Package: revtorch
RevTorchFramework for creating (partially) reversible neural networks with PyTorchRevTorch is introduced and explained in our paperA Partially Reversible U-Net for Memory-Efficient Volumetric Image Segmentation,
which was accepted for presentation atMICCAI 2019.If you find this code helpful in your research please cite the following paper:@article{PartiallyRevUnet2019Bruegger,
         author={Br{\"u}gger, Robin and Baumgartner, Christian F.
         and Konukoglu, Ender},
         title={A Partially Reversible U-Net for Memory-Efficient Volumetric Image Segmentation},
         journal={arXiv:1906.06148},
         year={2019},InstallationUse pip to install RevTorch:$pipinstallrevtorchRevTorch 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 thePyTorch website.UsageThis example shows how to use the RevTorch framework.importtorchimporttorchvisionimporttorchvision.transformsastransformsimporttorch.nnasnnimporttorch.nn.functionalasFimporttorch.optimasoptimimportrevtorchasrvdeftrain():trainset=torchvision.datasets.CIFAR10(root="./data",train=True,download=True,transform=transforms.ToTensor())trainloader=torch.utils.data.DataLoader(trainset,batch_size=4,shuffle=True)net=PartiallyReversibleNet()criterion=nn.CrossEntropyLoss()optimizer=optim.Adam(net.parameters())forepochinrange(2):running_loss=0.0fori,datainenumerate(trainloader,0):inputs,labels=dataoptimizer.zero_grad()outputs=net(inputs)loss=criterion(outputs,labels)loss.backward()optimizer.step()#logging stuffrunning_loss+=loss.item()LOG_INTERVAL=200ifi%LOG_INTERVAL==(LOG_INTERVAL-1):# print every 2000 mini-batchesprint('[%d,%5d] loss:%.3f'%(epoch+1,i+1,running_loss/LOG_INTERVAL))running_loss=0.0classPartiallyReversibleNet(nn.Module):def__init__(self):super(PartiallyReversibleNet,self).__init__()#initial non-reversible convolution to get to 32 channelsself.conv1=nn.Conv2d(3,32,3)#construct reversible sequencce with 4 reversible blocksblocks=[]foriinrange(4):#f and g must both be a nn.Module whos output has the same shape as its inputf_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 functionsblocks.append(rv.ReversibleBlock(f_func,g_func))#pack all reversible blocks into a reversible sequenceself.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)defforward(self,x):x=self.conv1(x)#the reversible sequence can be used like any other nn.Module. Memory-saving backpropagation is used automaticallyx=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])returnxif__name__=="__main__":train()Python versionTested with Python 3.6 and PyTorch 1.1.0. Should work with any version of Python 3.ChangelogVersion 0.2.4Added option to disable eager discarding of variables to allow for multiple backward() callsVersion 0.2.3Added option to use the same random seed for the forward and backwar pass (Pull request)Version 0.2.1Added option to select the dimension along which the tensor is split (Pull request)Version 0.2.0Fixed memory leak when not consuming output of the reversible block (Issue)Version 0.1.0Initial release