Hey @ClashLuke, nice code! You point to iRevNet as a comparison, and explain why this implementation is more memory efficient. However, iRevNet includes a method to, given the pre-pooler feature set for the CNN, directly invert the features and compute the input that generated those features. I don't see an obvious way to do that using your library. Is that something that can be done and do you intend to add a method to do that? I'm particularly interested in computing the inverse of a transformer given the output hidden states, which is something that iRevNet does not support as far as I can tell since it's CNN-specific, but if it's implementable within your framework as seems likely I would be very interested in leveraging it for my research.

Thank you! :)

Hi

I would like to use this library to build a ResNet20 model, I've tried several times but I still have the mismatched dimension error. My model is shown as follows:

import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.nn.init as init

from torch.nn.modules.batchnorm import _BatchNorm
from torch.nn.modules.instancenorm import _InstanceNorm

hidden_size = [16, 32, 64]

class View(nn.Module):
    def forward(self, x):
        batch_size = x.size(0)
        return x.view(batch_size, -1)

class LambdaLayer(nn.Module):
    def __init__(self, lambd):
        super(LambdaLayer, self).__init__()
        self.lambd = lambd

    def forward(self, x):
        return self.lambd(x)


class BasicBlock(nn.Module):
    expansion = 1

    def __init__(self, in_planes, planes, stride, norm_layer, conv_layer, option='A'):
        super(BasicBlock, self).__init__()
        
        self.bn1 = norm_layer(in_planes)
        self.conv1 = conv_layer(in_planes, planes, kernel_size=3, stride=stride, padding=1, bias=False)
        self.bn2 = norm_layer(planes)
        self.conv2 = conv_layer(planes, planes, kernel_size=3, stride=1, padding=1, bias=False)

        self.shortcut = nn.Sequential()
        if stride != 1 or in_planes != self.expansion * planes:
            if option == 'A':
                """
                For CIFAR10 ResNet paper uses option A.
                """
                self.shortcut = LambdaLayer(lambda x:
                                            F.pad(x[:, :, ::2, ::2], (0, 0, 0, 0, planes//4, planes//4), "constant", 0))
            elif option == 'B':
                self.shortcut = conv_layer(in_planes, self.expansion * planes, kernel_size=1, stride=stride, bias=False)

    def forward(self, x):
        out = F.relu(self.bn1(x))
        shortcut = self.shortcut(out) if hasattr(self, 'shortcut') else x
        print(1, shortcut.size())
        out = self.conv1(out)
        out = self.conv2(F.relu(self.bn2(out)))
        print(2, out.size())
        out += shortcut
        return out

class ResNet(nn.Module):

    def __init__(self, hidden_size, block, num_blocks, num_classes=10, bn_type='bn',
                 share_affine=False, track_running_stats=True):
        super(ResNet, self).__init__()
        self.in_planes = 16

        self.bn_type = bn_type
        if bn_type == 'bn':
            norm_layer = lambda n_ch: nn.BatchNorm2d(n_ch, track_running_stats=track_running_stats)
        elif bn_type == 'gn':
            norm_layer = lambda n_ch: nn.GroupNorm(4, n_ch) # 3 can be changed -- # of groups
        else:
            raise RuntimeError(f"Not support bn_type={bn_type}")
        conv_layer = nn.Conv2d
        first = conv_layer(3, hidden_size[0], kernel_size=3, stride=1, padding=1, bias=False)
        layer1 = self._make_layer(block, hidden_size[0], num_blocks[0], stride=1,
                                       norm_layer=norm_layer, conv_layer=conv_layer)
        layer2 = self._make_layer(block, hidden_size[1], num_blocks[1], stride=2,
                                       norm_layer=norm_layer, conv_layer=conv_layer)
        layer3 = self._make_layer(block, hidden_size[2], num_blocks[2], stride=2,
                                       norm_layer=norm_layer, conv_layer=conv_layer)
        
        self.rev_layers = revlib.ReversibleSequential(*[layer1, layer2, layer3])
        
        norm = norm_layer(hidden_size[2] * block.expansion)
        linear = nn.Linear(hidden_size[2] * block.expansion, num_classes)
        
        self.full_model = nn.Sequential(first, self.rev_layers, nn.ReLU(), norm, \
                                        nn.AdaptiveAvgPool2d((None, 1)), View(), linear)

    def _make_layer(self, block, planes, num_blocks, stride, norm_layer, conv_layer):
        strides = [stride] + [1] * (num_blocks - 1)
        layers = []
        for stride in strides:
            layers.append(block(self.in_planes, planes, stride, norm_layer, conv_layer))
            self.in_planes = planes * block.expansion
        return nn.Sequential(*layers)

    def forward(self, x):
        out = self.full_model(x)
        return out

def init_param(m):
    """Special init for ResNet"""
    if isinstance(m, (_BatchNorm, _InstanceNorm)):
        m.weight.data.fill_(1)
        m.bias.data.zero_()
    elif isinstance(m, nn.Linear):
        m.bias.data.zero_()
    return m

def resnet20(**kwargs):
    model = ResNet(hidden_size, BasicBlock, [3,3,3], **kwargs)
    model.apply(init_param)
    return model

I've tried to modify self.in_planes = 8 and hidden_size = [8, 16, 32], respectively, but it still does not work. Could you provide any hints? Is it possible to build a model in a forward way instead of wrapping reversible model with non-reversible layers like model = nn.Sequential(conv, rev_layer, conv)? I appreciate your help.

Hi!

Thank you for your great work. I was wondering how can I mode's information such as number of parameters, estimated total model size via torchsummary library or other possible tools.

I've received RuntimeError: Given groups=1, weight of size [256, 64, 3, 3], expected input[2, 32, 224, 224] to have 64 channels, but got 32 channels instead when I run summary(rev_model, input_size=(64, 224, 224)). I understand that the input in the revnet is divided into two parts in terms of channels (64/2=32). Any suggestion? I appreciate your help!

If I need to stride the input, I'm wondering how you can recommend how to use it.

class Stride(nn.Module):
    def forward(self, input):
        return input[[slice(None)]*2 + [slice(0, None, 2) for _ in range(2, input.ndim)]]
Stride()(torch.randn(1, 1, 11, 16)).shape # 1, 1, 6, 8

PyTorch 1.10 will add torch.autograd.graph.saved_tensors_hooks, which I used to significantly improve code style without hurting performance or changing the API.

I used the improved code to add a memory_savings=True switch to ReversibleSequential without further complicating the code. In addition, the flag allows for easier debugging and tests.

This branch also adds tests regarding the memory efficiency and correctness of gradients.
All tests pass on both branches.