tor

torchprofile

A general and accurate MACs / FLOPs profiler for PyTorch models

Showing:

Popularity

Downloads/wk

0

GitHub Stars

321

Maintenance

Last Commit

2mos ago

Contributors

4

Package

Dependencies

3

License

MIT

Categories

Readme

Torchprofile

This is a profiler to count the number of MACs / FLOPs of PyTorch models based on torch.jit.trace.

  • It is more general than ONNX-based profilers as some operations in PyTorch are not supported by ONNX for now.
  • It is more accurate than hook-based profilers as they cannot profile operations within torch.nn.Module.

Installation

pip install torchprofile

Getting Started

You should first define your PyTorch model and its (dummy) input:

import torch
from torchvision.models import resnet18

model = resnet18()
inputs = torch.randn(1, 3, 224, 224)

You can then measure the number of MACs using profile_macs:

from torchprofile import profile_macs

macs = profile_macs(model, inputs)

License

This repository is released under the MIT license. See LICENSE for additional details.

Rate & Review

Great Documentation0
Easy to Use0
Performant0
Highly Customizable0
Bleeding Edge0
Responsive Maintainers0
Poor Documentation0
Hard to Use0
Slow0
Buggy0
Abandoned0
Unwelcoming Community0
100