TorchProfile

https://github.com/zhijian-liu/torchprofile

Admin

TorchProfile counts the number of multiply-accumulate operations (MACs) in PyTorch models by tracing...

Tokens:556

Snippets:9

Trust Score:6.8

License:MIT

Update:1 month ago

Show doc for...

Tokens:

Raw

### Install torchprofile with pip

Source: https://github.com/zhijian-liu/torchprofile/blob/main/README.md

Install the torchprofile package from PyPI.

```bash
pip install torchprofile
```

--------------------------------

### Per-operator MAC breakdown with reduction=None

Source: https://github.com/zhijian-liu/torchprofile/blob/main/README.md

Get a per-operator breakdown of MACs by passing reduction=None. The results dictionary maps nodes to their MAC counts; only nodes with positive MACs are printed, showing scope, operator, and MMACs.

```python
results = profile_macs(model, inputs, reduction=None)
for node, macs in results.items():
    if macs > 0:
        print(f"{node.scope:40s} {node.operator:30s} {macs / 1e6:>8.2f} MMACs")
```

--------------------------------

### Profile MACs of a PyTorch model

Source: https://github.com/zhijian-liu/torchprofile/blob/main/README.md

Profile a Hugging Face model with torchprofile. The model is loaded from the meta-llama/Llama-3.2-1B checkpoint and evaluated on a random input tensor of shape (1, 128). The total MACs are printed in GMACs.

```python
import torch
from transformers import AutoModel
from torchprofile import profile_macs

model = AutoModel.from_pretrained("meta-llama/Llama-3.2-1B").eval()
inputs = torch.randint(0, model.config.vocab_size, (1, 128))

macs = profile_macs(model, inputs)
print(f"{macs / 1e9:.2f} GMACs")
```

Tokens:

Raw