Package: memcnn
APyTorchframework 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.htmlFeaturesEnable memory savings during training by wrapping arbitrary invertible PyTorch functions with theInvertibleModuleWrapperclass.Simple toggling of memory saving by setting thekeep_inputproperty of theInvertibleModuleWrapper.Turn arbitrary non-linear PyTorch functions into invertible versions using theAdditiveCouplingor theAffineCouplingclasses.Training and evaluation code for reproducing RevNet experiments using MemCNN.CI tests for Python v3.7 and torch v1.0, v1.1, v1.4 and v1.7 with good code coverage.ExamplesCreating an AdditiveCoupling with memory savingsimporttorchimporttorch.nnasnnimportmemcnn# define a new torch Module with a sequence of operations: Relu o BatchNorm2d o Conv2dclassExampleOperation(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))defforward(self,x):returnself.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 waymodel_normal=ExampleOperation(channels=10)model_normal.eval()Y=model_normal(X)# turn the ExampleOperation invertible using an additive couplinginvertible_module=memcnn.AdditiveCoupling(Fm=ExampleOperation(channels=10//2),Gm=ExampleOperation(channels=10//2))# test that it is actually a valid invertible module (has a valid inverse method)assertmemcnn.is_invertible_module(invertible_module,test_input_shape=X.shape)# wrap our invertible_module using the InvertibleModuleWrapper and benefit from memory savings during traininginvertible_module_wrapper=memcnn.InvertibleModuleWrapper(fn=invertible_module,keep_input=True,keep_input_inverse=True)# by default the module is set to training, the following sets this to evaluation# note that this is required to pass input tensors to the model with requires_grad=False (inference only)invertible_module_wrapper.eval()# test that the wrapped module is also a valid invertible moduleassertmemcnn.is_invertible_module(invertible_module_wrapper,test_input_shape=X.shape)# compute the forward pass using the wrapperY2=invertible_module_wrapper.forward(X)# the input (X) can be approximated (X2) by applying the inverse method of the wrapper on Y2X2=invertible_module_wrapper.inverse(Y2)# test that the input and approximation are similarasserttorch.allclose(X,X2,atol=1e-06)Run PyTorch ExperimentsAfter installing MemCNN run:python-mmemcnn.train[MODEL][DATASET][--fresh][--no-cuda]Available values forDATASETarecifar10andcifar100.Available values forMODELareresnet32,resnet110,resnet164,revnet38,revnet110,revnet164Use the--freshflag to remove earlier experiment results.Use the--no-cudaflag 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 thepythondirective with the appropriate Python 3 directive. For example when using the MemCNN docker image usepython3.6.When MemCNN was installed usingpipor from sources you might need to setup a configuration file before running this command.
Read the corresponding section about how to do this here:https://memcnn.readthedocs.io/en/latest/installation.html