Package: invtorch
InvTorch: Memory-Efficient Invertible FunctionsWhen working with extremely deep neural networks, memory becomes an immediate concern. It will be filled withsaved tensorsthat are needed for gradient computation. PyTorch provides a solution,checkpoint_sequential, that allows us to segment every few layers as a checkpoint. Such that, the forward pass is run inno_gradmode. Meanwhile, the inputs of every segment is saved in memory. Every other unreferenced tensor gets deallocated. In the backward pass, the forward pass of every segment, starting from the last to the first, will be run again using its saved inputs to compute its gradients. Refer tothisfor more details.This module extends the functionality oftorch.utils.checkpoint.checkpointto work with invertible functions. So, everything now can be released from memory and recomputed later using the inverse function in the backward pass. This is useful for extremely wide networks where more compute is traded with memory. However, there are few considerations to keep in mind when working with invertible checkpoints and non-materialized tensors. Please, read the limitations section below for specifics.InstallationInvTorchhas minimal dependencies. It only requires Python>=3.6and PyTorch>=1.10.0.pipinstallinvtorchUsageWe are interested ininvtorch.nn.Modulewhich inherits fromtorch.nn.Module. Subclass it to implement your own invertible code.importtorchfromtorchimportnnimportinvtorch.nnasinnclassInvertibleLinear(inn.Module):def__init__(self,in_features,out_features):super().__init__()self.weight=nn.Parameter(torch.randn(out_features,in_features))self.bias=nn.Parameter(torch.randn(out_features))deffunction(self,inputs):returninputs@self.weight.T+self.biasdefinverse(self,outputs):return(outputs-self.bias)@self.weight.T.pinverse()if__name__=='__main__':x=torch.randn(10,3)model=InvertibleLinear(3,5)print('Is invertible:',model.check(x))y=model(x)print('Input was freed:',x.storage().size()==0)y.backward(torch.randn_like(y))print('Input was restored:',x.storage().size()!=0)forward()You can immediately notice few differences to the regular PyTorch module here. There is no longer a need to defineforward(). Instead, it is replaced withfunction(). Additionally, it is necessary to define its inverse function asinverse(). Both methods should input and output only positional arguments as atupleor a singletorch.Tensor. Furthermore,forward()accepts an optional keyword argumentkeepwhich is an iterable of input tensors that shouldn't be deallocated.function()The first call tofunction()is always run indry_mode(check it withinvtorch.in_dry_mode()). This is a novel mode that has gradient graph construction enabled but doesn't allow for backward propagation. The second call is during the backward pass which is when the gradients will actually be computed (check it withinvtorch.in_backward()).inverse()You can verify your implementation ofinverse()by callingcheck(). In some cases, switching to double precision is advised as invertible functions can run into some numerical instability when using single precision. For some functions, a view of an input tensor is passed in the outputs. In such case, this will be automatically detected and the input tensor will not be released from memory.num_outputsSometimes,inverse()requires auxiliary outputs fromfunction()that are not necessarily needed as an output forforward(). To specify how many outputs do we need to keep, we can modify thenum_outputsproperty (Nonekeeps everything). Extra outputs will be hidden usingprocess()which is automatically called byforward(). Another propertynum_inputsis also defined which helps composing inverse calls like ininvtorch.nn.Sequential.inverse().reverse()invtorch.nn.Modulecan be implemented to be reversible, i.e.forward()will callinverse()instead offunction(). Not all invertible modules need to support reversibility. If you want to support it on your own module, then you need to override thereversibleproperty to returnTrue. The module can be reversed by callingreverse()and checked with thereversedproperty. To avoid confusion,Modulehascall_function()andcall_inverse()which will call the correct function based on the value ofreversed.checkpoint, invertible, and seedinvtorch.nn.Modulehas two flags which control the mode of operation;checkpointandinvertible. Ifcheckpointwas set toFalse, or when working inno_gradmode, it acts exactly as a normal PyTorch module. Otherwise, the model is eitherinvertibleor acheckpointdepending on whetherinvertibleis set toTrueorFalse, respectively. Those, flags can be changed at any time during operation without any repercussions. A third flagseedis by defaultFalseand if set toTrue, it ensures that the forward runs in the same random number generator's state of the devices of the input tensors.invtorch.checkpoint()PyTorch's checkpoint cannot track therequires_gradattribute for its output tensors since it is running inno_gradmode. Instead, InvTorch's checkpoint doesn't have this issue because it runs indry_mode. In addition, it supports invertible functions.LimitationsThere are few caveats to consider though. Invertible functions are hard to define without requiring more memory. Moreover, they are prone to numerical instabilities (e.g., multiplying by numbers close to zero). Even if we can get away with these fundamental problems, there are technical details to consider. There is no way of guarding against accessing the data in the input tensors after callingfunction()and before the backward pass. It is up to the user to ensure this. Otherwise, it is possible to run into illegal memory access errors. Think of residual connections as an example. Inx + f(x), assumingfis an invertible checkpoint,xwill be freed from memory before the sum is computed. On the other hand, we can maybe usex.clone() + f(x)(notf(x) + x.clone()!) but now we have a copy ofxin memory. It is recommended to encapsulate this insidefitself or use a simple checkpoint instead. Other alternatives exists and you should study your case carefully before deciding to use this. For instance, check outtorch.autograd.graph.saved_tensors_hooks()andtorch.autograd.graph.save_on_cpu().