Could it support models with multiple input values in the "forward" method?

[fourson]

Dear author:
I have encountered such a situation that my model have multiple input values with different input_shape. And it seems that this library can only support the models with only a single input value. Could you please update the library to support models with multiple input values?
Thanks!

[frgfm]

Hello @fourson 👋

Thanks for your suggestion! I thought about this and I think it's doable, however:

• for now, this library only records PyTorch modules operations (i.e. if you write a forward method that has an extra pooling operations (no params), that won't be hooked and thus won't be recorded as such)
• with that in mind, I can modify the library so that it works with PyTorch layers that respect that. Meaning that if you use a transformer layer for instance, that's quite alright. If you do something very custom where a complex operation is implemented with the functional API, Torchscan will be missing a lot.

So could you share a short snippet to show the type of model you are referring to please? :)
(when I do a heavy modification, I like to collect expectations from users to make sure I can maximize support for everyone)

Cheers ✌️

[fourson]

Hello @frgfm, thanks for your reply!
Here I provide my code snippet:

def forward(self, E, B, Bi):
        B_edge_feature = self.B_edge_feature_extraction(torch_laplacian(B))
        E_feature = self.E_feature_extraction(E)
        motion_clues = self.fuse1(torch.cat((B_edge_feature, E_feature), dim=1))
        F = self.flow_block(motion_clues)

        if self.rgb:
            motion_clues = self.up(motion_clues)

        Bi_gamma = Bi ** (1 / 2.2)
        Bi_clean = torch.clamp(self.Bi_denoise(Bi_gamma) + Bi_gamma, min=0, max=1) ** 2.2
        Bi_feature = self.Bi_feature_extraction(Bi_clean)
        backbone_out = self.backbone(self.fuse2(torch.cat((Bi_feature, motion_clues), dim=1)))
        S = torch.clamp(self.out_block(backbone_out) + Bi_clean, min=0, max=1)
        return F, Bi_clean, S

where torch_laplacian is only a simple custom function:

@torch.jit.script
def torch_laplacian(img_tensor):
    # (N, C, H, W) image tensor -> (N, C, H, W) edge tensor, the same as cv2.Laplacian
    padded = F.pad(img_tensor, pad=[1, 1, 1, 1], mode='reflect')
    return padded[:, :, 2:, 1:-1] + padded[:, :, 0:-2, 1:-1] + padded[:, :, 1:-1, 2:] + padded[:, :, 1:-1, 0:-2] - \
           4 * img_tensor

The input args in my ``forward'' method do not share the same shape:
E: (N, C1, H1,W1)
B: (N, C2, H2, W2)
Bi: (N, C3, H3, W3)
where N denotes the batch size

and the methods starting with self. is just the layers defined in __init__().

I think it would be very helpful if your library could provide an API that support ``custom input shape''.
\eg,

summary_custom(model, input_shapes, max_depth=2)

where input_shapes is a dict mapping the args in the forward method to their actual shape:

input_shapes = {
'E': (C1, H1,W1),
'B': (C2, H2,W2),
'Bi': (C3, H3,W3)
}

or you could directly let the users to build the input tensors instead of generate a random one inside the library.
\eg,

summary_custom(model, input_tensors, max_depth=2)

where input_tensors is a dict mapping the args in the forward method to their actual tensor:

input_shapes = {
'E': tensor_E,
'B': tensor_B,
'Bi': tensor_Bi
}

[frgfm]

I think it's interesting. My reasoning now is:

1. if this makes its way to torchscan, the same function will be used, but instead of passing a tuple, you'll be passing a list of tuples
2. now there are two challenges:
   • for each module, computing in a standard way all the relevant statistics
   • transcribe this into a readable format

With the latter, I see two problems:

• if we allow multiple inputs, we have to allow multiple outputs. Fair enough but that requires some work
• second and most importantly: torchscan has predone work, meaning that I wrote functions for every single supported module types to compute everything. Now, in your custom forward method: if the "custom" part changes the number of FLOPs etc compared to what is already being done, this won't work. Otherwise, I'd need to have written functions for any possible custom forward functions

To solve the last part, there is only a single way to solve this: trace the model. That will uniformly convert modules & functions to ONNX for instance, which will allow this. With pure PyTorch, there is no hooks on functions, so it's limited to modules.

What do you think? :)