How to count macs and parameters during forwarding in models or layers by pytorch?

user_16849546 0
Viewed 747
class model(nn.Module)
    def __init__()
        self.l1 = nn.linear(10,10)
        self.l2 = nn.linear(10,10)
    def forward(x)
        y1 = self.l1(x)
        y2 = self.l2(x)
        y3 = y1 + y2
        return y2 + y2 + y3*y1

In this example, I could use forward_hook functions to trace two linear layers and their parameters.fn is hook function.

m.register_forward_hook(fn)

However, y3 is not counted as a parameter and the macs of y2 + y2 + y3*y1 is not counted in macs, too.
How can I solve this?

"macs" is a way of measuring layers' complexity. For example, y1 *(y2 + y3) is one macs, if y1, y2, y3 are floats.

pytorch deep-learning machine-learning complexity-theory python
1 Answers

Seems to me like PyTorch-OpCounter is a perfect fit for your use case:

  1. Define your nn.Module class:

    class NN(nn.Module):
    def __init__(self):
        super().__init__()
        self.l1 = nn.Linear(10, 10)
        self.l2 = nn.Linear(10, 10)

    def forward(self, x):
        y1 = self.l1(x)
        y2 = self.l2(x)
        y3 = y1 + y2
        return y2 + y2 + y3*y1

  2. Initialize your model and input:

    model = NN()
x = torch.empty(16, 10)

  3. Import thop and call the profiler:

    from thop import profile
macs, params = profile(model, inputs=(x,))

There you have your results:

>>> macs, params
(3200.0, 220.0)
user_63313690
Related