Spaces:
Running
Running
File size: 2,807 Bytes
a80d6bb c74a070 a80d6bb c74a070 a80d6bb c74a070 |
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 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 |
import torch
from fvcore.nn import FlopCountAnalysis
from einops.einops import rearrange
from src import get_model_cfg
from src.models.backbone import FPN as topicfm_featnet
from src.models.modules import TopicFormer
from src.utils.dataset import read_scannet_gray
from third_party.loftr.src.loftr.utils.cvpr_ds_config import default_cfg
from third_party.loftr.src.loftr.backbone import ResNetFPN_8_2 as loftr_featnet
from third_party.loftr.src.loftr.loftr_module import LocalFeatureTransformer
def feat_net_flops(feat_net, config, input):
model = feat_net(config)
model.eval()
flops = FlopCountAnalysis(model, input)
feat_c, _ = model(input)
return feat_c, flops.total() / 1e9
def coarse_model_flops(coarse_model, config, inputs):
model = coarse_model(config)
model.eval()
flops = FlopCountAnalysis(model, inputs)
return flops.total() / 1e9
if __name__ == "__main__":
path_img0 = "assets/scannet_sample_images/scene0711_00_frame-001680.jpg"
path_img1 = "assets/scannet_sample_images/scene0711_00_frame-001995.jpg"
img0, img1 = read_scannet_gray(path_img0), read_scannet_gray(path_img1)
img0, img1 = img0.unsqueeze(0), img1.unsqueeze(0)
# LoFTR
loftr_conf = dict(default_cfg)
feat_c0, loftr_featnet_flops0 = feat_net_flops(
loftr_featnet, loftr_conf["resnetfpn"], img0
)
feat_c1, loftr_featnet_flops1 = feat_net_flops(
loftr_featnet, loftr_conf["resnetfpn"], img1
)
print(
"FLOPs of feature extraction in LoFTR: {} GFLOPs".format(
(loftr_featnet_flops0 + loftr_featnet_flops1) / 2
)
)
feat_c0 = rearrange(feat_c0, "n c h w -> n (h w) c")
feat_c1 = rearrange(feat_c1, "n c h w -> n (h w) c")
loftr_coarse_model_flops = coarse_model_flops(
LocalFeatureTransformer, loftr_conf["coarse"], (feat_c0, feat_c1)
)
print(
"FLOPs of coarse matching model in LoFTR: {} GFLOPs".format(
loftr_coarse_model_flops
)
)
# TopicFM
topicfm_conf = get_model_cfg()
feat_c0, topicfm_featnet_flops0 = feat_net_flops(
topicfm_featnet, topicfm_conf["fpn"], img0
)
feat_c1, topicfm_featnet_flops1 = feat_net_flops(
topicfm_featnet, topicfm_conf["fpn"], img1
)
print(
"FLOPs of feature extraction in TopicFM: {} GFLOPs".format(
(topicfm_featnet_flops0 + topicfm_featnet_flops1) / 2
)
)
feat_c0 = rearrange(feat_c0, "n c h w -> n (h w) c")
feat_c1 = rearrange(feat_c1, "n c h w -> n (h w) c")
topicfm_coarse_model_flops = coarse_model_flops(
TopicFormer, topicfm_conf["coarse"], (feat_c0, feat_c1)
)
print(
"FLOPs of coarse matching model in TopicFM: {} GFLOPs".format(
topicfm_coarse_model_flops
)
)
|