ultralytics-thop

import torch import torch.nn as nn from thop import profile, clever_format # Define a simple PyTorch model class SimpleModel(nn.Module): def __init__(self): super(SimpleModel, self).__init__() self.conv1 = nn.Conv2d(3, 64, kernel_size=3, padding=1) self.relu = nn.ReLU() self.pool = nn.MaxPool2d(kernel_size=2, stride=2) self.conv2 = nn.Conv2d(64, 128, kernel_size=3, padding=1) self.fc = nn.Linear(128 * 8 * 8, 10) # Assuming input size 3x32x32 initially def forward(self, x): x = self.pool(self.relu(self.conv1(x))) x = self.pool(self.relu(self.conv2(x))) x = torch.flatten(x, 1) x = self.fc(x) return x # Instantiate the model and create a dummy input model = SimpleModel() input_tensor = torch.randn(1, 3, 32, 32) # Profile the model macs, params = profile(model, inputs=(input_tensor,), verbose=False) # Format the output for readability macs_formatted, params_formatted = clever_format([macs, params], "%.3f") print(f"Model MACs: {macs_formatted}") print(f"Model Params: {params_formatted}")