Fix calculate_macs() for Linear layers.

[andravin]

Fixes #77.

[TylerYep]

Could you add a unit test and also explain why this fixes the issue?

Hardcoding Linear is not ideal, but if there is something special about Linear layers that requires this fix I would like to understand it more completely.

[andravin]

The documentation for torch.nn.Linear says the input tensor can have "any number of dimensions."

Currently, torchinfo only handles two-dimensional input tensors correctly, any additional dimensions are ignored:

                    self.macs += self.output_size[0] * cur_params

The change accounts for all of the leading dimensions:

                elif "Linear" in self.class_name:
                    self.macs += int(cur_params * prod(self.output_size[:-1]))

You can think of torch.nn.Linear as a tensor-matrix multiplication or n-mode product, with the last dimension of the tensor as the "mode." This is equivalent to matrix multiplication if we first unfold all of the leading dimensions of the tensor into a single dimension. The size of the unfolded dimension is prod(self.output_size[:-1]).

See section A.2.3 N-mode Product in this reference for a detailed explanation.

[andravin]

I corrected the output files for the LSTM tests, but the flan_t5 tests are still broken. I am not familiar with that model,so I would have to learn how to calculate the number of operations correctly.

[andravin]

The way LayerInfo.calculate_macs just checks a partial match on the class-name for Conv and now Linear and ignores the module name seems dicey. Exact matches against fully qualified <module>.<class> names would seem more robust.

Also, any unrecognized layer is assumed to be a Linear layer with 2D input, judging by the code in the else block. That also seems error-prone. It might be better to throw an exception for unsupported layers.

[andravin]

I added a unit test for torch.nn.Linear that uses a 3D input tensor.

[TylerYep]

Thanks for the explanation. It seems very reasonable that this error would cause other model outputs to change, so feel free to update those as suggested by the test errors.

I think this specific calculation has a lot of room for improvement, since it was only hardcoded with a few layer types to begin with. It would be great to enhance the <module>.<class> exact match code as well, but feel free to tackle these things in a separate PR.

[andravin]

It looks like torch.nn.Embedding layers have the same bug that torch.nn.Linear layers did.

[andravin]

@TylerYep, I followed your guidance and changed the ground truth in flan_t5_small.out to equal the output of the unit tests with the linear layer MACs fix.

Let me know if you prefer the three commits in this pull request to be squashed into a single commit.

[TylerYep]

Looks good. I'll happily accept more PRs expanding this functionality. Thank you for your contributions!