You signed in with another tab or window. Reload to refresh your session.You signed out in another tab or window. Reload to refresh your session.You switched accounts on another tab or window. Reload to refresh your session.Dismiss alert
!pip install -q thop
import torch
from utils import DEVICE
from thop import profile, clever_format
# Display the DEVICE
print(f"DEVICE: {DEVICE}")
# Assuming 'model' is your PyTorch model and 'input' is a tensor representing the input to the model
input = torch.randn(1, 10) # example input for a single length 10 sequnce
input = input.to(DEVICE)
model = torch.nn.Linear(10, 11) # example model
model = model.to(DEVICE)
# model.eval()
# Measure the FLOPs
macs, params = profile(model, inputs=(input,), verbose=False)
print(f"MACs: {macs}, params: {params}")
macs, params = clever_format([macs, params], "%.3f")
print(f"MACs: {macs}, params: {params}")
num_params = sum(p.numel() for p in model.parameters() if p.requires_grad)
print(f"Parameters: {num_params}")
Outputs:
For some reason the resnet example from the README seems to work ok.
Output:
Why wouldn't this work for such basic torch.nn.Modules like Linear and LSTM but work for more complicated models?
The text was updated successfully, but these errors were encountered: