Lyken17 / pytorch-OpCounter

Count the MACs / FLOPs of your PyTorch model.
MIT License
4.85k stars 528 forks source link

Add support for nn.Transformer modules. #138

Closed HaoKang-Timmy closed 3 years ago

HaoKang-Timmy commented 3 years ago

Add a hook function to count the macs of nn.Transformer.