text: from collections import deque
  from typing import List, Mapping, Tuple
  import torch
  import torch.nn as nn
  import torch.nn.functional as F
  import math
  
  from torch.nn.modules import loss
  from lib.models.decode import _nms, _topk
  from lib.models.utils import _tranpose_and_gather_feat, _sigmoid
  from lib.models.losses import FocalLoss
  from threading import Lock
  
  
  def print_meta(out):
      if isinstance(out, torch.Tensor):
          return 'Tensor{%s}' % str(out.shape)
      if isinstance(out, List):
          return [print_meta(x) for x in out]
      if isinstance(out, Tuple):
          return tuple([print_meta(x) for x in out])
      if isinstance(out, Mapping):
          return {k: print_meta(v) for k, v in out.items()}
      return out
  
  
  class CmpRelMod(nn.Module):
  
      def __init__(self,
                   in_channels=1024,
                   feat_dim=1024,
                   groups=16,
                   stages=3,
                   use_pos=False,
                   use_gt=False,
                   lvl_sep=False,
                   num_lvls=5,
                   finest_scale=56,
                   ignore_min=0,
                   add_type='sum',
                   norm=True,
                   qlt_mask=True,
                   qlt_type='cmp',
                   loss_weight=1.0,
                   **kwargs
                   ):
          super(CmpRelMod, self).__init__()
          self.groups = groups
          self.group = groups
          self.stages = stages
          self.feat_dim = feat_dim
          self.in_channels = in_channels
          self.use_pos = use_pos
          self.use_gt = use_gt
          self.lvl_sep = lvl_sep
          self.num_lvls = num_lvls
          self.finest_scale = finest_scale
          self.ignore_min = ignore_min
          self.add_type = add_type
          self.norm = norm
          self.qlt_mask = qlt_mask
          self.qlt_type = qlt_type
          self.loss_weight = loss_weight
          self.embed_dim = self.feat_dim // self.groups
          assert self.stages > 0
          self.fcs = []
          self.kws = []
          self.qws = []
          self.vws = []
          self.qlts = []
          self.bns = []
          for i in range(self.stages):
              if i == 0:
                  self.fcs.append(nn.Linear(self.in_channels, self.feat_dim))
              else:
                  self.fcs.append(nn.Linear(self.feat_dim, self.feat_dim))
              self.kws.append(nn.Linear(self.feat_dim, self.feat_dim))
              self.qws.append(nn.Linear(self.feat_dim, self.feat_dim))
              self.qlts.append(nn.Sequential(
                  nn.Linear(self.feat_dim, self.feat_dim),
                  nn.ReLU(),
                  nn.BatchNorm1d(self.feat_dim),
                  nn.Linear(self.feat_dim, 1),
                  # nn.ReLU(),
              ))
              self.vws.append(nn.Conv2d(self.feat_dim * self.groups, self.feat_dim,
                                        kernel_size=1, stride=1, padding=0, groups=self.groups))
              setattr(self, 'us_%d' % i, nn.Parameter(torch.Tensor(
                  self.groups, 1, self.embed_dim)))
              self.bns.append(nn.BatchNorm1d(self.feat_dim))
          if self.use_pos:
              self.wgs = []
              for i in range(self.stages):
                  self.wgs.append(
                      nn.Conv2d(self.embed_dim, self.groups, kernel_size=1, stride=1, padding=0))
              self.wgs = nn.ModuleList(self.wgs)
          self.fcs = nn.ModuleList(self.fcs)
          self.kws = nn.ModuleList(self.kws)
          self.qws = nn.ModuleList(self.qws)
          self.vws = nn.ModuleList(self.vws)
          # self.us = nn.ParameterList(self.us)
          self.qlts = nn.ModuleList(self.qlts)
          self.bns = nn.ModuleList(self.bns)
          if self.feat_dim != self.in_channels:
              self.adjust = nn.Linear(self.feat_dim, self.in_channels)
          if self.add_type == 'group':
              self.aff_alpha = nn.Parameter(torch.Tensor(self.groups, 2))
          if self.add_type == 'cond':
              self.aff_alpha = nn.Linear(self.feat_dim, self.groups * 2)
          self.init_weights()
  
      def init_weights(self):
          fcs = [self.kws, self.qws]
          for em in fcs:  # linear layers
              for ei in em:
                  nn.init.kaiming_uniform_(ei.weight, a=1)
                  nn.init.constant_(ei.bias, 0)
          convs = [self.vws]
          if self.use_pos:
              convs.append(self.wgs)
          for em in convs:  # conv layers
              for ei in em:
                  nn.init.normal_(ei.weight, std=0.01)
                  nn.init.constant_(ei.bias, 0)
          if self.feat_dim != self.in_channels:
              nn.init.normal_(self.adjust.weight, std=0.01)
              nn.init.constant_(self.adjust.bias, 0)
          for u in range(self.stages):  # single parameter
              nn.init.normal_(getattr(self, 'us_%d' % u), std=0.01)
  
          for em in self.qlts:
              for ei in em:
                  if isinstance(ei, nn.Linear):
                      nn.init.kaiming_uniform_(ei.weight, a=1)
                      nn.init.constant_(ei.bias, 0)
  
          if self.add_type == 'group':
              for p in [self.aff_alpha]:
                  nn.init.normal_(p, std=0.01)
          if self.add_type == 'cond':
              for ei in [self.aff_alpha]:
                  nn.init.kaiming_uniform_(ei.weight, a=1)
                  nn.init.constant_(ei.bias, 0)
  
      def feat_aff(self, feats, ref_feats, qw, kw, us, position_embedding=None, gw=None, outter_aff=None, add='avg', mask=None):
  
          if position_embedding is not None:
              assert gw is not None
              position_feat_1 = F.relu(gw(position_embedding))
              # aff_weight, [num_rois, group, num_nongt_rois, 1]
              aff_weight = position_feat_1.permute(2, 1, 3, 0)
              # aff_weight, [num_rois, group, num_nongt_rois]
              aff_weight = aff_weight.squeeze(3)
  
          def ignore_part(aff_softmax_):
              aff_flatten: torch.Tensor = aff_softmax_.reshape(
                  -1, aff_softmax_.shape[-1])
              ref_num = aff_flatten.size(1)
              ref_num_dist = max(int(ref_num * (1 - self.ignore_min)), 1)
              scores, inds = aff_flatten.topk(ref_num_dist, dim=1)
              pinds = torch.arange(aff_flatten.size(0)).to(
                  aff_flatten.device) * ref_num
              pinds = pinds.unsqueeze(-1) + inds
              pinds = pinds.reshape(-1)
              mask = aff_flatten.new_zeros(aff_flatten.shape).reshape(-1)
              mask[pinds] = 1
              mask = mask.reshape(*aff_softmax_.shape)
              max_e, _ = aff_softmax_.max(dim=2)
              max_e = max_e.unsqueeze(2)
              exp_s = (aff_softmax_ - max_e).exp()
              chosen_ = mask * exp_s
              eps = 1e-19
              ss = 1. / (chosen_.sum(dim=2).unsqueeze(2) + eps)
              aff_softmax_ = chosen_ * ss
              # aff_softmax_ = F.softmax(aff_softmax_, dim=2) * mask
              # k = aff_softmax_.sum(dim=2).unsqueeze(-1)
              # s = 1. / k
              # aff_softmax_ = aff_softmax_ * s
              return aff_softmax_
  
          q_feats = qw(feats)
          q_feats = q_feats.reshape(-1, self.groups,
                                    self.feat_dim // self.groups)
          q_feats = q_feats.permute(1, 0, 2)
          k_feats = kw(ref_feats)
          k_feats = k_feats.reshape(-1, self.groups,
                                    self.feat_dim // self.groups)
          k_feats = k_feats.permute(1, 0, 2)
          aff_a = torch.bmm(q_feats, k_feats.transpose(1, 2))
          aff_c = torch.bmm(us, k_feats.transpose(1, 2))
          aff = aff_a + aff_c
  
          aff_scale = (1.0 / math.sqrt(float(self.feat_dim / self.groups))) * aff
          # aff_scale, [num_rois, group, num_nongt_rois]
          aff_scale = aff_scale.permute(1, 0, 2)
  
          if position_embedding is not None:
              weighted_aff = (aff_weight + 1e-6).log() + aff_scale
          else:
              weighted_aff = aff_scale
          raw_aff = weighted_aff
  
          if outter_aff is not None and add != 'inner_only':
              if add == 'avg':
  
                  aff_softmax = F.softmax(weighted_aff, dim=2)
                  outter_aff_softmax = F.softmax(outter_aff, dim=2)
  
                  aff_softmax_ = (
                      outter_aff_softmax + aff_softmax) / 2.
  
              elif add == 'sum':
                  aff = weighted_aff + outter_aff
  
                  if self.ignore_min > 0:
                      aff_softmax_ = ignore_part(aff)
                  else:
                      aff_softmax_ = F.softmax(aff, dim=2)
  
              elif add == 'max':
                  aff_softmax_ = torch.max(weighted_aff, outter_aff)
                  aff_softmax_ = F.softmax(aff_softmax_, dim=2)
              elif add == 'group':
                  aff_alpha = self.aff_alpha.softmax(dim=1)
                  aff_softmax_ = weighted_aff * \
                      aff_alpha[:, 0].reshape(
                          1, -1, 1) + outter_aff * aff_alpha[:, 1].reshape(1, -1, 1)
                  if self.ignore_min > 0:
                      aff_softmax_ = ignore_part(aff_softmax_)
                  else:
                      aff_softmax_ = F.softmax(aff_softmax_, dim=2)
              elif add == 'cond':
                  aff_alpha = self.aff_alpha(feats).reshape(-1, self.groups, 2)
                  aff_alpha = aff_alpha.softmax(dim=2)
                  aff_softmax_ = weighted_aff * \
                      aff_alpha[:, :, 0].unsqueeze(
                          2) + outter_aff * aff_alpha[:, :, 1].unsqueeze(2)
                  aff_softmax_ = F.softmax(aff_softmax_, dim=2)
          else:
              aff_softmax_ = weighted_aff
              if self.norm:
                  max_e, _ = aff_softmax_.max(dim=2)
                  max_e = max_e.unsqueeze(2)
                  exp_s = (aff_softmax_ - max_e).exp()
              if self.ignore_min > 0:
                  aff_softmax_ = ignore_part(aff_softmax_)
              else:
                  aff_softmax_ = F.softmax(aff_softmax_, dim=2)
  
          if mask is not None:
              if self.norm:
                  chosen_ = (mask[:, None, :] * exp_s)
                  # all_sum = exp_s.sum(dim=2)
                  eps = 1e-19
                  ss = 1. / (chosen_.sum(dim=2).unsqueeze(2) + eps)
                  aff_softmax_ = chosen_ * ss
                  k = 1
                  aff_softmax_ = (aff_softmax_ ** k) / ((aff_softmax_ ** k).sum(dim=2).unsqueeze(2) + eps)
              else:
                  aff_softmax_ = aff_softmax_ * mask[:, None, :]
          # print(aff_softmax_.shape, 'pred', aff_softmax_[:3, :1].sum(dim=2))
          aff_softmax_reshape = aff_softmax_.reshape(
              aff_softmax_.shape[0] * aff_softmax_.shape[1], aff_softmax_.shape[2])
          return aff_softmax_reshape, raw_aff
  
      @staticmethod
      def extract_position_embedding(position_mat, feat_dim, wave_length=1000.0):
          device = position_mat.device
          # position_mat, [num_rois, num_nongt_rois, 4]
          feat_range = torch.arange(0, feat_dim / 8, device=device)
          dim_mat = torch.full((len(feat_range),), wave_length, device=device).pow(
              8.0 / feat_dim * feat_range)
          dim_mat = dim_mat.view(1, 1, 1, -1).expand(*position_mat.shape, -1)
  
          position_mat = position_mat.unsqueeze(
              3).expand(-1, -1, -1, dim_mat.shape[3])
          position_mat = position_mat * 100.0
  
          div_mat = position_mat / dim_mat
          sin_mat, cos_mat = div_mat.sin(), div_mat.cos()
  
          # [num_rois, num_nongt_rois, 4, feat_dim / 4]
          embedding = torch.cat([sin_mat, cos_mat], dim=3)
          # [num_rois, num_nongt_rois, feat_dim]
          embedding = embedding.reshape(
              embedding.shape[0], embedding.shape[1], embedding.shape[2] * embedding.shape[3])
  
          return embedding
  
      @staticmethod
      def extract_position_matrix(bbox, ref_bbox):
          xmin, ymin, xmax, ymax = torch.chunk(ref_bbox, 4, dim=1)
          bbox_width_ref = xmax - xmin + 1
          bbox_height_ref = ymax - ymin + 1
          center_x_ref = 0.5 * (xmin + xmax)
          center_y_ref = 0.5 * (ymin + ymax)
  
          xmin, ymin, xmax, ymax = torch.chunk(bbox, 4, dim=1)
          bbox_width = xmax - xmin + 1
          bbox_height = ymax - ymin + 1
          center_x = 0.5 * (xmin + xmax)
          center_y = 0.5 * (ymin + ymax)
  
          delta_x = center_x - center_x_ref.transpose(0, 1)
          delta_x = delta_x / bbox_width
          delta_x = (delta_x.abs() + 1e-3).log()
  
          delta_y = center_y - center_y_ref.transpose(0, 1)
          delta_y = delta_y / bbox_height
          delta_y = (delta_y.abs() + 1e-3).log()
  
          delta_width = bbox_width / bbox_width_ref.transpose(0, 1)
          delta_width = delta_width.log()
  
          delta_height = bbox_height / bbox_height_ref.transpose(0, 1)
          delta_height = delta_height.log()
  
          position_matrix = torch.stack(
              [delta_x, delta_y, delta_width, delta_height], dim=2)
  
          return position_matrix
  
      def cal_pos_embed(self, b_rois, a_rois):
          m = self.extract_position_matrix(b_rois[:, 1:], a_rois[:, 1:])
          pos_embed = self.extract_position_embedding(m, feat_dim=64)
          pos_embed = pos_embed.permute(2, 0, 1)
          # [1, 64, num_rois, num_nongt_rois]
          pos_embed = pos_embed.unsqueeze(0)
  
          return pos_embed
  
      def map_roi_levels(self, rois):
          """Map rois to corresponding feature levels by scales.
  
          - scale < finest_scale * 2: level 0
          - finest_scale * 2 <= scale < finest_scale * 4: level 1
          - finest_scale * 4 <= scale < finest_scale * 8: level 2
          - scale >= finest_scale * 8: level 3
  
          Args:
              rois (Tensor): Input RoIs, shape (k, 5).
  
          Returns:
              Tensor: Level index (0-based) of each RoI, shape (k, )
          """
          scale = torch.sqrt(
              (rois[:, 3] - rois[:, 1] + 1) * (rois[:, 4] - rois[:, 2] + 1))
          target_lvls = torch.floor(torch.log2(scale / self.finest_scale + 1e-6))
          target_lvls = target_lvls.clamp(min=0, max=self.num_levels - 1).long()
          return target_lvls
  
      def level_exclusion(self, roi_x, roi_y):
          lvl_x = self.map_roi_levels(roi_x)
          lvl_y = self.map_roi_levels(roi_y)
          ret = lvl_x.view(-1, 1) == lvl_y.view(1, -1)
          ret = ret.float()
          return ret
  
      def forward(self, cur, ref, gt=None, return_ref=False):
          x_cur, rois_cur, scores_cur = cur
          x_ref, rois_ref, scores_ref = ref
          affs = []
          sims = []
          qualities = []
          feats = []
          for i in range(self.stages):
              gw = self.wgs[i] if self.use_pos else None
              x_cur = F.relu(self.fcs[i](x_cur))
              x_ref = F.relu(self.fcs[i](x_ref))
              pos_emb1 = None
              pos_emb2 = None
              if self.use_pos:
                  pos_emb1 = self.cal_pos_embed(rois_cur, rois_ref)
                  pos_emb2 = self.cal_pos_embed(rois_ref, rois_ref)
              quality1 = self.qlts[i](x_cur.detach())
              quality2 = self.qlts[i](x_ref.detach())
              mask1 = None
              mask2 = None
              if self.qlt_mask:
                  if self.qlt_type == 'cmp':
                      mask1 = (quality1.view(-1, 1) <
                               quality2.view(1, -1)).float()
                      mask2 = (quality2.view(-1, 1) <
                               quality2.view(1, -1)).float()
                  else:
                      assert isinstance(self.qlt_type, (int, float))
                      mask1 = (quality2.view(1, -1) > self.qlt_type).float()
                      mask1 = mask1.repeat(quality1.size(0), 1)
                      mask2 = (quality2.view(1, -1) > self.qlt_type).float()
                      mask2 = mask2.repeat(quality2.size(0), 1)
              # qmask = mask1
              if self.lvl_sep:
                  mask1 = mask1 * self.level_exclusion(rois_cur, rois_ref)
                  mask2 = mask2 * self.level_exclusion(rois_ref, rois_ref)
              aff_1, aff_sim = self.feat_aff(
                  x_cur, x_ref, self.qws[i], self.kws[i], getattr(self, 'us_%d' % i), pos_emb1, gw, add=self.add_type, mask=mask1)
              aff_2, _ = self.feat_aff(
                  x_ref, x_ref, self.qws[i], self.kws[i], getattr(self, 'us_%d' % i), pos_emb2, gw, add=self.add_type, mask=mask2)
              affs.append(aff_1)
              sims.append(aff_sim)
              qualities.append(quality1)
              feats.append(x_cur)
              if self.use_gt:
                  assert gt is not None
                  gt_process, kwargs = gt
                  if gt_process is not None:
                      # print(aff_sim[:10, :, :20], 'sim')
                      # print(mask1[:10, :20], 'mask')
                      aff_1, aff_2 = gt_process(x_cur, x_ref, rois_cur, rois_ref, index=i, aff1=aff_1.reshape(
                          -1, self.group, aff_1.size(1)), aff2=aff_2.reshape(-1, self.group, aff_2.size(1)), **kwargs)
                      if aff_1.size(1) != self.group:
                          aff_1 = aff_1.repeat(1, self.group, 1)
                          aff_2 = aff_2.repeat(1, self.group, 1)
                      aff_1 = aff_1.reshape(-1, aff_1.size(2))
                      aff_2 = aff_2.reshape(-1, aff_2.size(2))
                      # TODO: gt aff
              # output_t, [num_rois * group, feat_dim]
              output_cur = torch.matmul(aff_1, x_ref)
              # output_t, [num_rois, group * feat_dim, 1, 1]
              output_cur = output_cur.reshape(-1,
                                              self.group * self.feat_dim, 1, 1)
              # linear_out, [num_rois, dim[2], 1, 1]
              linear_cur = self.vws[i](output_cur)
  
              x_cur = x_cur + linear_cur.squeeze(3).squeeze(2)
              x_cur = self.bns[i](x_cur)
  
              if i != self.stages - 1 and not return_ref:
                  # output_t, [num_rois * group, feat_dim]
                  output_ref = torch.matmul(aff_2, x_ref)
                  # output_t, [num_rois, group * feat_dim, 1, 1]
                  output_ref = output_ref.reshape(-1,
                                                  self.group * self.feat_dim, 1, 1)
                  # linear_out, [num_rois, dim[2], 1, 1]
                  linear_ref = self.vws[i](output_ref)
  
                  x_ref = x_ref + linear_ref.squeeze(3).squeeze(2)
                  x_ref = self.bns[i](x_ref)
          if self.feat_dim != self.in_channels:
              x_cur = self.adjust(x_cur)
              if return_ref:
                  x_ref = self.adjust(x_ref)
          if return_ref:
              return x_cur, affs, sims, qualities, feats, x_ref
          return x_cur, affs, sims, qualities, feats
  
  
  class SQGARelMod(nn.Module):
  
      def __init__(self, stages=2, use_gt=0) -> None:
          super(SQGARelMod, self).__init__()
          self.rel = CmpRelMod(
              in_channels=256, feat_dim=256, groups=1, stages=stages, use_gt=bool(use_gt))
          self.pre_cls = nn.Sequential(
              nn.Linear(256, 256),
              nn.ReLU(),
              nn.Linear(256, 1),
          )
          nn.init.normal_(self.pre_cls[0].weight, std=0.1)
          nn.init.normal_(self.pre_cls[2].weight, std=0.1)
          self.binary_loss = FocalLoss()
          # bank_size = 7
          # self.feature_bank = deque(maxlen=bank_size)
          # self.label_bank = deque(maxlen=bank_size)
          # self.per_ref_num = 256
          # self.lock = Lock()
  
      def forward(self, inputs, inputs_ref, gt=None):
          # print(print_meta(inputs))
          # hm = _nms(inputs[0]['hm'])
          # scores, inds, clses, ys, xs = _topk(hm, K=self.per_ref_num)
          # # print(inds.shape)
          # raw_feats = inputs[0]['raw']
          # selected_feats = _tranpose_and_gather_feat(raw_feats, inds)
          # batch_size, dim, h, w = raw_feats.shape
          # reshaped_feats = selected_feats.reshape(-1, dim)
          # # for one in self.feature_bank:
          # #     print(one.device)
          # self.lock.acquire()
          # ref_feats = torch.cat(list(map(lambda x: x.to(
          #     reshaped_feats.device), self.feature_bank)) + [reshaped_feats])
          # ref_labels = list(
          #     map(lambda x: x.to(reshaped_feats.device), self.label_bank))
          # self.lock.release()
          # # print(ref_feats.device)
          # raw_feats = raw_feats.permute(0, 2, 3, 1).reshape(-1, dim)
          # steps = (torch.arange(batch_size, device=raw_feats.device, dtype=torch.long)
          #          * (h * w)).reshape(batch_size, 1).repeat(1, self.per_ref_num)
          # inds = (steps + inds).reshape(-1)
          # # print(inds.shape, inds.unique().shape)
          # cur_feats = raw_feats[inds]
          # cur_feats = inputs
          # ref_feats = torch.cat(inputs_ref)
          # print(cur_feats.shape, cur_feats.device)
          # print(ref_feats.shape, cur_feats.device)
          refined_feats, affs, sims, qlts, sfeats = self.rel(
              inputs, inputs_ref, gt=gt)
          # print(refined_feats.shape)
          # raw_feats[inds] = refined_feats
          # raw_feats = raw_feats.reshape(
          #     batch_size, h, w, dim).permute(0, 3, 1, 2)
          # # print(raw_feats.shape)
          # outputs = {
          #     'raw': raw_feats,
          #     'id': inputs[0]['id']
          # }
          return refined_feats, (affs, sims, qlts, sfeats, inputs[0])
  
      def loss(self, stuff, cur_labels, ref_labels):
          affs, sims, qlts, sfeats, cur_feats = stuff
          # print(cur_labels.shape, ref_labels.shape)
          losses = {}
          # print(cur_labels[:10])
          # print(ref_labels[:20])
          sim_target = cur_labels.view(-1, 1) == ref_labels.view(1, -1)
          # print(sim_target.shape)
          # print(sims[0].shape)
          weight_1 = sim_target.new_zeros(sim_target.shape).int()
          weight_2 = sim_target.new_zeros(sim_target.shape).int()
          pos_cur = torch.nonzero(cur_labels > 0)
          pos_ref = torch.nonzero(ref_labels > 0)
          weight_1[pos_cur, :] = 1
          weight_2[:, pos_ref] = 1
          weight = (weight_1 + weight_2).flatten()
          pos_pos_inds = torch.nonzero(weight == 2)
          pos_neg_inds = torch.nonzero(weight == 1)
          neg_neg_inds = torch.nonzero(weight == 0)
          pos_pos_cnt = pos_pos_inds.size(0)
          pos_neg_cnt = pos_pos_cnt * 2
          neg_neg_cnt = pos_pos_cnt * 1
          rand_inds = torch.randperm(pos_neg_inds.size(0))
          rand_inds = pos_neg_inds[rand_inds[:pos_neg_cnt]]
          weight[rand_inds] = 2
          rand_inds = torch.randperm(neg_neg_inds.size(0))
          rand_inds = neg_neg_inds[rand_inds[:neg_neg_cnt]]
          weight[rand_inds] = 2
  
          # weight[:10, :10] = 1
          # weight[pos_cur, :10] = 1
          # weight[:10, pos_ref] = 1
          sim_inds = torch.nonzero(weight == 2)
          # print(sim_inds.shape)
          # print(sim_inds.shape, pos_cur.numel(), pos_ref.numel())
          sim_target = sim_target.reshape(-1).float()
          # print(sim_target.shape)
          # print(sim_inds.shape, weight.shape, sim_target.shape)
          sim_target = sim_target[sim_inds]
          # print(sim_inds[:10], sim_target.shape)
          # print(sims[0].flatten()[:10], sim_target[:10])
          # print(torch.nonzero(sim_target > 0).size(0), torch.nonzero(sim_target == 0).size(0), pos_pos_cnt, weight.numel())
          # print(self.rel.stages)
          for i in range(self.rel.stages):
              sim = sims[i].sum(dim=1).reshape(-1)[sim_inds]
              sim = _sigmoid(sim)
              qlt = _sigmoid(qlts[i])
              sfeat = sfeats[i]
              # print(sfeat.device)
              # print(self.pre_cls[0].weight.device)
              pre_cls_pred = _sigmoid(self.pre_cls(sfeat))
              # print(pre_cls_pred[:10])
              # print(cls_target[:10])
              losses['pre_cls.%d.loss' % i] = self.binary_loss(
                  pre_cls_pred.flatten(), cur_labels.flatten().float()) / self.rel.stages * 0.5
              qlt_target = (1 - torch.abs(cur_labels.flatten().float() - pre_cls_pred.flatten())).detach()
              # print(qlt_target[:10])
              losses['qlt.%d.loss' % i] = self.binary_loss(
                  qlt.flatten(), qlt_target.flatten()) / self.rel.stages * 0.5
              losses['sim.%d.loss' % i] = self.binary_loss(
                  sim.flatten(), sim_target.flatten()) / self.rel.stages
          # print(losses['sim.%d.loss' % 0])
          # input()
          loss = 0
          for v in losses.values():
              loss += v
          loss = loss * 0.5
          # losses['loss_rel'] = loss
          # for k in list(losses.keys()):
          #     losses[k] = float(losses[k])
          return loss, losses
  