TorchSummary for revlib

[taokz]

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!

[ClashLuke]

Good find! RevLib currently splits the input into two parts for you, so you'd have to give it double the number of features you want.
Do you know how we could make this clearer from the library side?

[taokz]

Thank you for your reply.

Do you mean doubling number of channels? such as summary(rev_model, input_size=(128, 224, 224)) if I set the model as follows?

channels = 64
channel_multiplier = 4
depth = 3
classes = 1000
# Create a reversible model. f() is invoked depth-times with different weights.
rev_model = revlib.ReversibleSequential(*[block() for _ in range(depth)])

However, when I modify that, I've received the other error:

File ~\Anaconda3\lib\site-packages\torchsummary\torchsummary.py:19, in summary.<locals>.register_hook.<locals>.hook(module, input, output)
     17 m_key = "%s-%i" % (class_name, module_idx + 1)
     18 summary[m_key] = OrderedDict()
---> 19 summary[m_key]["input_shape"] = list(input[0].size())
     20 summary[m_key]["input_shape"][0] = batch_size
     21 if isinstance(output, (list, tuple)):

AttributeError: 'tuple' object has no attribute 'size'

[ClashLuke]

It seems like TorchSummary can't handle tuple outputs very well. Looking through their code, I also found that TorchSummary won't count in any weight not stored in .weight or .bias (like in many custom models) in their total parameter count. To fix these issues, I submitted a PR to TorchSummary. For now, it might be best to install my fork directly via python3 -m pip install --upgrade git+https://github.com/ClashLuke/pytorch-summary.

I get an output like this:

>>> import torch
>>> import torchsummary
>>> import revlib
>>> block = lambda: torch.nn.Linear(10, 10)
>>> r = revlib.ReversibleSequential(*[block() for _ in range(4)])
>>> r.cuda()
ReversibleSequential(
  (stem): Sequential(
    (0): ReversibleModule(
      coupling_forward=additive_coupling_forward
      coupling_inverse=additive_coupling_inverse
      target_device=None
      (wrapped_module): ReversibleWrapper(
        (wrapped_module): Linear(in_features=10, out_features=10, bias=True)
      )
    )
    (1): ReversibleModule(
      coupling_forward=additive_coupling_forward
      coupling_inverse=additive_coupling_inverse
      target_device=None
      (wrapped_module): ReversibleWrapper(
        (wrapped_module): Linear(in_features=10, out_features=10, bias=True)
      )
    )
    (2): ReversibleModule(
      coupling_forward=additive_coupling_forward
      coupling_inverse=additive_coupling_inverse
      target_device=None
      (wrapped_module): ReversibleWrapper(
        (wrapped_module): Linear(in_features=10, out_features=10, bias=True)
      )
    )
    (3): ReversibleModule(
      coupling_forward=additive_coupling_forward
      coupling_inverse=additive_coupling_inverse
      target_device=None
      (wrapped_module): ReversibleWrapper(
        (wrapped_module): Linear(in_features=10, out_features=10, bias=True)
      )
    )
  )
)
>>> torchsummary.summary(r, input_size=(4, 10))
----------------------------------------------------------------
        Layer (type)               Output Shape         Param #
================================================================
            Linear-1              ([2, 2, 10],)             110
 ReversibleWrapper-2   ([2, 2, 10], [2, 2, 10])               0
  ReversibleModule-3  (([2, 2, 10], [2, 2, 10], [2, 2, 10], [2, 2, 10]),)               0
            Linear-4              ([2, 2, 10],)             110
 ReversibleWrapper-5   ([2, 2, 10], [2, 2, 10])               0
  ReversibleModule-6  (([2, 2, 10], [2, 2, 10], [2, 2, 10], [2, 2, 10]),)               0
            Linear-7              ([2, 2, 10],)             110
 ReversibleWrapper-8   ([2, 2, 10], [2, 2, 10])               0
  ReversibleModule-9  (([2, 2, 10], [2, 2, 10], [2, 2, 10], [2, 2, 10]),)               0
           Linear-10              ([2, 2, 10],)             110
ReversibleWrapper-11   ([2, 2, 10], [2, 2, 10])               0
 ReversibleModule-12  (([2, 2, 10], [2, 2, 10], [2, 2, 10], [2, 2, 10]),)               0
ReversibleSequential-13              ([2, 4, 10],)               0
================================================================
Total params: 440
Trainable params: 440
Non-trainable params: 0
----------------------------------------------------------------
Input size (MB): 0.00
Forward/backward pass size (MB): 78.18
Params size (MB): 0.00
Estimated Total Size (MB): 78.18
----------------------------------------------------------------

(440, 440)
>>>

[taokz]

Awesome! Thanks!