File size: 687 Bytes
e970115 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 |
import argparse
import torch
from fvcore.nn import FlopCountAnalysis
from ultralytics import YOLO
if __name__ == '__main__':
parser = argparse.ArgumentParser()
parser.add_argument('model', type=str, help='Model path for validation.')
parser.add_argument('--imgsz', default=928, type=int, help='Image size to validate.')
args = parser.parse_args()
model = torch.load(args.model, map_location='cpu')['model'].float()
fca = FlopCountAnalysis(
model=model.eval(),
inputs=torch.rand(1, 3, args.imgsz, args.imgsz),
)
fca.unsupported_ops_warnings(False)
fca.uncalled_modules_warnings(False)
print(f"{fca.total() * 1e-9:.2f} GMACS ")
|