Batch mode slow about lightglue HOT 4 CLOSED

      n = 30 

      config={'width_confidence': 0.99,
              'depth_confidence': 0.95}
      
      Model_lg=LightGlue(pretrained='superpoint', **config).eval()
      local_feature_matcher = Model_lg.to('cuda')

      keys = ['descriptors0', 'keypoints0', 'image_size0', 'keypoint_scores0',
              'descriptors1', 'keypoints1', 'image_size1', 'keypoint_scores1']
      
      for key in keys:
          pred[key] = torch.cat([pred[key]] * n)
      
      
      for i in range(10):
          last_time=time()
          # for i in range(30):
          with torch.inference_mode():
              pred1 = local_feature_matcher(pred)
          current_time=time()
          print(current_time-last_time)

        n = 1
        
        keys = ['descriptors0', 'keypoints0', 'image_size0', 'keypoint_scores0',
                'descriptors1', 'keypoints1', 'image_size1', 'keypoint_scores1']
        
        for key in keys:
            pred[key] = torch.cat([pred[key]] * n)
        
        
        for i in range(10):
            last_time=time()
            for i in range(30):
                with torch.inference_mode():
                    pred1 = local_feature_matcher(pred)
            current_time=time()
            print(current_time-last_time)

      def apply_cached_rotary_emb(
              freqs: torch.Tensor, t: torch.Tensor) -> torch.Tensor:
          t_rot=rotate_half(t).permute(1,0,2,3)
          t=t.permute(1,0,2,3)
          return ((t * freqs[0]) + (t_rot * freqs[1])).permute(1,0,2,3)

      def _forward(self, data: dict) -> dict:
          for key in self.required_data_keys:
              assert key in data, f'Missing key {key} in data'
          kpts0_, kpts1_ = data['keypoints0'], data['keypoints1']
          b, m, _ = kpts0_.shape
          b, n, _ = kpts1_.shape
  
          kpts0 = normalize_keypoints(
              kpts0_, size=data.get('image_size0'), shape=None)
          kpts1 = normalize_keypoints(
              kpts1_, size=data.get('image_size1'), shape=None)
  
          assert torch.all(kpts0 >= -1) and torch.all(kpts0 <= 1)
          assert torch.all(kpts1 >= -1) and torch.all(kpts1 <= 1)
  
          desc0 = data['descriptors0'].detach().permute(0,2,1)
          desc1 = data['descriptors1'].detach().permute(0,2,1)
  
          assert(desc0.shape[-1] == self.conf.input_dim)
          assert(desc1.shape[-1] == self.conf.input_dim)
  
          if torch.is_autocast_enabled():
              desc0 = desc0.half()
              desc1 = desc1.half()
  
          desc0 = self.input_proj(desc0)
          desc1 = self.input_proj(desc1)
  
          # cache positional embeddings
          encoding0 = self.posenc(kpts0)
          encoding1 = self.posenc(kpts1)
  
          # GNN + final_proj + assignment
          ind0 = torch.arange(0, m).to(device=kpts0.device).expand(b,-1)
          ind1 = torch.arange(0, n).to(device=kpts0.device).expand(b,-1)
          prune0 = torch.ones_like(ind0)  # store layer where pruning is detected
          prune1 = torch.ones_like(ind1)
          dec, wic = self.conf.depth_confidence, self.conf.width_confidence
          token0, token1 = None, None
          for i in range(self.conf.n_layers):
              # self+cross attention
              desc0, desc1 = self.self_attn[i](
                  desc0, desc1, encoding0, encoding1)
              desc0, desc1 = self.cross_attn[i](desc0, desc1)
              if i == self.conf.n_layers - 1:
                  continue  # no early stopping or adaptive width at last layer
              if dec > 0:  # early stopping
                  token0, token1 = self.token_confidence[i](desc0, desc1)
                  if self.stop(token0, token1, self.conf_th(i), dec, m+n):
                      break
              if wic > 0:  # point pruning
                  match0, match1 = self.log_assignment[i].scores(desc0, desc1)
                  mask0 = self.get_mask(token0, match0, self.conf_th(i), 1-wic)
                  mask1 = self.get_mask(token1, match1, self.conf_th(i), 1-wic)
                  ind0, ind1 = ind0[mask0][None], ind1[mask1][None]
                  desc0, desc1 = desc0[mask0][None], desc1[mask1][None]
                  if desc0.shape[-2] == 0 or desc1.shape[-2] == 0:
                      break
                  encoding0 = encoding0[:, :, mask0][:, None]
                  encoding1 = encoding1[:, :, mask1][:, None]
              prune0[:, ind0] += 1
              prune1[:, ind1] += 1
  
          if wic > 0:  # scatter with indices after pruning
              scores_, _ = self.log_assignment[i](desc0, desc1)
              dt, dev = scores_.dtype, scores_.device
              scores = torch.zeros(b, m+1, n+1, dtype=dt, device=dev)
              scores[:, :-1, :-1] = -torch.inf
              scores[:, ind0[0], -1] = scores_[:, :-1, -1]
              scores[:, -1, ind1[0]] = scores_[:, -1, :-1]
              x, y = torch.meshgrid(ind0[0], ind1[0], indexing='ij')
              scores[:, x, y] = scores_[:, :-1, :-1]
          else:
              scores, _ = self.log_assignment[i](desc0, desc1)
  
          m0, m1, mscores0, mscores1 = filter_matches(
              scores, self.conf.filter_threshold)
  
          return {
              'log_assignment': scores,
              'matches0': m0,
              'matches1': m1,
              'matching_scores0': mscores0,
              'matching_scores1': mscores1,
              'stop': i+1,
              'prune0': prune0,
              'prune1': prune1,
          }