Spaces:
Running
Running
| 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 | |
| ) | |
| ) | |