Spaces:
No application file
No application file
File size: 1,372 Bytes
b26e93d |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 |
"""
Copyright (c) 2024 The D-FINE Authors. All Rights Reserved.
"""
import os
import sys
sys.path.insert(0, os.path.join(os.path.dirname(os.path.abspath(__file__)), "../.."))
import argparse
import torch
import torch.nn as nn
from calflops import calculate_flops
from src.core import YAMLConfig
def custom_repr(self):
return f"{{Tensor:{tuple(self.shape)}}} {original_repr(self)}"
original_repr = torch.Tensor.__repr__
torch.Tensor.__repr__ = custom_repr
def main(
args,
):
"""main"""
cfg = YAMLConfig(args.config, resume=None)
class Model_for_flops(nn.Module):
def __init__(
self,
) -> None:
super().__init__()
self.model = cfg.model.deploy()
def forward(self, images):
outputs = self.model(images)
return outputs
model = Model_for_flops().eval()
flops, macs, _ = calculate_flops(
model=model, input_shape=(1, 3, 640, 640), output_as_string=True, output_precision=4
)
params = sum(p.numel() for p in model.parameters())
print("Model FLOPs:%s MACs:%s Params:%s \n" % (flops, macs, params))
if __name__ == "__main__":
parser = argparse.ArgumentParser()
parser.add_argument(
"--config", "-c", default="configs/dfine/dfine_hgnetv2_l_coco.yml", type=str
)
args = parser.parse_args()
main(args)
|