def forward(self, outputs, targets): """ Performs the matching Args: outputs (dict): This is a dict that contains at least these entries: "pred_logits": Tensor of dim [batch_size, num_queries, num_classes] with the classification logits "pred_boxes": Tensor of dim [batch_size, num_queries, 4] with the predicted box coordinates targets (list): This is a list of targets (len(targets) = batch_size), where each target is a dict: "labels": Tensor of dim [num_target_boxes] (where num_target_boxes is the number of ground-truth objects in the target) containing the class labels "boxes": Tensor of dim [num_target_boxes, 4] containing the target box coordinates Returns: A list of size batch_size, containing tuples of (index_i, index_j) where: - index_i is the indices of the selected predictions (in order) - index_j is the indices of the corresponding selected targets (in order) For each batch element, it holds: len(index_i) = len(index_j) = min(num_queries, num_target_boxes) """ bs, num_queries = outputs["pred_logits"].shape[:2] # We flatten to compute the cost matrices in a batch out_prob = outputs["pred_logits"].flatten(0, 1).softmax( -1) # [batch_size * num_queries, num_classes] out_bbox = outputs["pred_boxes"].flatten( 0, 1) # [batch_size * num_queries, 4] # Also concat the target labels and boxes tgt_ids = torch.cat([v["labels"] for v in targets]) tgt_bbox = torch.cat([v["boxes"] for v in targets]) # Compute the classification cost. Contrary to the loss, we don't use the NLL, # but approximate it in 1 - proba[target class]. # The 1 is a constant that doesn't change the matching, it can be ommitted. cost_class = -out_prob[:, tgt_ids] # Compute the L1 cost between boxes cost_bbox = torch.cdist(out_bbox, tgt_bbox, p=1) # Compute the giou cost betwen boxes cost_giou = -generalized_box_iou(box_cxcywh_to_xyxy(out_bbox), box_cxcywh_to_xyxy(tgt_bbox)) # Final cost matrix C = self.cost_bbox * cost_bbox + self.cost_class * \ cost_class + self.cost_giou * cost_giou C = C.view(bs, num_queries, -1).cpu() sizes = [len(v["boxes"]) for v in targets] indices = [ linear_sum_assignment(c[i]) for i, c in enumerate(C.split(sizes, -1)) ] return [(torch.as_tensor(i, dtype=torch.int64), torch.as_tensor(j, dtype=torch.int64)) for i, j in indices]
def get_distances(representations, sigma=1): rview = representations.view(representations.size(0),-1) distances = torch.cdist(rview,rview,p=2) return distances
def forward(self, outputs_dict, targets): """ Performs the matching Returns: A list of size batch_size, containing tuples of (index_i, index_j) where: - index_i is the indices of the selected predictions (in order) - index_j is the indices of the corresponding selected targets (in order) For each batch element, it holds: len(index_i) = len(index_j) = min(num_queries, num_target_boxes) """ outputs = outputs_dict['pred_det'] bs, num_queries = outputs["pred_logits"].shape[:2] # We flatten to compute the cost matrices in a batch out_prob = outputs["pred_logits"].flatten(0, 1).softmax(-1) # [batch_size * num_queries, num_classes] out_bbox = outputs["pred_boxes"].flatten(0, 1) # [batch_size * num_queries, 4] # Also concat the target labels and boxes tgt_ids = torch.cat([v["labels"] for v in targets]) tgt_bbox = torch.cat([v["boxes"] for v in targets]) # Compute the classification cost. Contrary to the loss, we don't use the NLL, # but approximate it in 1 - proba[target class]. # The 1 is a constant that doesn't change the matching, it can be ommitted. cost_class = -out_prob[:, tgt_ids] # Compute the L1 cost between boxes cost_bbox = torch.cdist(out_bbox, tgt_bbox, p=1) # Compute the giou cost betwen boxes cost_giou = -generalized_box_iou(box_cxcywh_to_xyxy(out_bbox), box_cxcywh_to_xyxy(tgt_bbox)) # Final cost matrix C = self.cost_bbox * cost_bbox + self.cost_class * cost_class + self.cost_giou * cost_giou C = C.view(bs, num_queries, -1).cpu() sizes = [len(v["boxes"]) for v in targets] indices = [linear_sum_assignment(c[i]) for i, c in enumerate(C.split(sizes, -1))] indices = [(torch.as_tensor(i, dtype=torch.int64), torch.as_tensor(j, dtype=torch.int64)) for i, j in indices] if outputs_dict['pred_rel'] is None: indices_dict = { 'det': indices, 'rel': None } return indices_dict # for rel rel_outputs = outputs_dict['pred_rel'] bs, rel_num_queries = rel_outputs["pred_logits"].shape[:2] rel_out_prob = rel_outputs["pred_logits"].flatten(0, 1).sigmoid() # [batch_size * num_queries, num_classes] rel_out_bbox = rel_outputs["pred_boxes"].flatten(0, 1) # [batch_size * num_queries, 4] rel_tgt_ids = torch.cat([v["rel_labels"] for v in targets]) rel_tgt_bbox = torch.cat([v["rel_vecs"] for v in targets]) # interaction category semantic distance rel_cost_list = [] for idx, r_tgt_id in enumerate(rel_tgt_ids): tgt_rel_id = torch.where(r_tgt_id == 1)[0] rel_cost_list.append(-(rel_out_prob[:, tgt_rel_id]).sum( dim=-1) * self.cost_class) rel_cost_class = torch.stack(rel_cost_list, dim=-1) # another implementation # rel_cost_class = -(rel_out_prob * rel_tgt_ids).sum( # dim=-1) * self.cost_class) # interaction vector location distance rel_cost_bbox = torch.cdist(rel_out_bbox, rel_tgt_bbox, p=1) # Final cost matrix rel_C = self.cost_bbox * rel_cost_bbox + self.cost_class * rel_cost_class rel_C = rel_C.view(bs, rel_num_queries, -1).cpu() rel_sizes = [len(v["rel_vecs"]) for v in targets] rel_indices = [linear_sum_assignment(c[i]) for i, c in enumerate(rel_C.split(rel_sizes, -1))] rel_indices = [(torch.as_tensor(i, dtype=torch.int64), torch.as_tensor(j, dtype=torch.int64)) for i, j in rel_indices] indices_dict = { 'det': indices, 'rel': rel_indices, } return indices_dict
def pairwise_euaclidean_distance(x, y): return torch.cdist(x, y)
def euclidean(x, y): distances = torch.cdist(x, y, p=2.0, compute_mode="donot_use_mm_for_euclid_dist") return distances
all_distances.append(cur_assoc_distance.item()) if cur_assoc_distance <= CORRECT_THRESHOLD: correct_count += 1 all_correct_counts.append(correct_count) COUNT_CORRECT_greedy = True if COUNT_CORRECT_greedy: CORRECT_THRESHOLD = .15 all_pred_clusters = torch.stack(params, dim=1).squeeze()[:, :, :2] # print("len(params):", len(params)) # print("params[0].shape:", params[0].shape) # print("all_pred_clusters.shape:", all_pred_clusters.shape) # sleep(asldfhj) gt_objects = batch['gt_objects'].cuda() batch_size = all_pred_clusters.shape[0] pairwise_distances = torch.cdist(all_pred_clusters.contiguous(), gt_objects.contiguous()) (batch_size1, object_count, pred_count) = pairwise_distances.shape assert (batch_size1 == batch_size) # print("pairwise_distances.shape:", pairwise_distances.shape) correct_count = 0 distance_loss = 0.0 for i in range(object_count): #find min distances per batch values, indices = pairwise_distances.view(batch_size, -1).min(dim=1) # distance_loss = distance_loss + torch.sum(values) distance_loss = distance_loss + torch.sum(values) correct_count += torch.sum(values < CORRECT_THRESHOLD).item() # print("correct_count:", correct_count) row_indices = indices // pred_count row_indices = row_indices.unsqueeze(dim=1) #.unsqueeze(dim=1)
def forward(ctx, input, weight): ctx.save_for_backward(input, weight) out = -torch.cdist(input, weight, p=1) return out
def _calculate_pairwise_dist(self, X, U): return torch.cdist(X, X)
def forward(self, outputs, targets): """ Performs the matching Params: outputs: This is a dict that contains at least these entries: "pred_logits": Tensor of dim [batch_size, num_queries, num_classes] with the classification logits "pred_boxes": Tensor of dim [batch_size, num_queries, 4] with the predicted box coordinates targets: This is a list of targets (len(targets) = batch_size), where each target is a dict containing: "labels": Tensor of dim [num_target_boxes] (where num_target_boxes is the number of ground-truth objects in the target) containing the class labels "boxes": Tensor of dim [num_target_boxes, 4] containing the target box coordinates Returns: A list of size batch_size, containing tuples of (index_i, index_j) where: - index_i is the indices of the selected predictions (in order) - index_j is the indices of the corresponding selected targets (in order) For each batch element, it holds: len(index_i) = len(index_j) = min(num_queries, num_target_boxes) """ bs, num_queries = outputs["pred_logits"].shape[:2] # We flatten to compute the cost matrices in a batch if self.use_focal: out_prob = outputs["pred_logits"].flatten( 0, 1).sigmoid() # [batch_size * num_queries, num_classes] out_bbox = outputs["pred_boxes"].flatten( 0, 1) # [batch_size * num_queries, 4] else: out_prob = outputs["pred_logits"].flatten(0, 1).softmax( -1) # [batch_size * num_queries, num_classes] out_bbox = outputs["pred_boxes"].flatten( 0, 1) # [batch_size * num_queries, 4] # Also concat the target labels and boxes tgt_ids = torch.cat([v["labels"] for v in targets]) tgt_bbox = torch.cat([v["boxes_xyxy"] for v in targets]) # Compute the classification cost. Contrary to the loss, we don't use the NLL, # but approximate it in 1 - proba[target class]. # The 1 is a constant that doesn't change the matching, it can be ommitted. if self.use_focal: # Compute the classification cost. alpha = self.focal_loss_alpha gamma = self.focal_loss_gamma neg_cost_class = (1 - alpha) * (out_prob**gamma) * ( -(1 - out_prob + 1e-8).log()) pos_cost_class = alpha * ( (1 - out_prob)**gamma) * (-(out_prob + 1e-8).log()) cost_class = pos_cost_class[:, tgt_ids] - neg_cost_class[:, tgt_ids] else: cost_class = -out_prob[:, tgt_ids] # Compute the L1 cost between boxes image_size_out = torch.cat( [v["image_size_xyxy"].unsqueeze(0) for v in targets]) image_size_out = image_size_out.unsqueeze(1).repeat(1, num_queries, 1).flatten(0, 1) image_size_tgt = torch.cat([v["image_size_xyxy_tgt"] for v in targets]) out_bbox_ = out_bbox / image_size_out tgt_bbox_ = tgt_bbox / image_size_tgt cost_bbox = torch.cdist(out_bbox_, tgt_bbox_, p=1) # Compute the giou cost betwen boxes # cost_giou = -generalized_box_iou(box_cxcywh_to_xyxy(out_bbox), box_cxcywh_to_xyxy(tgt_bbox)) cost_giou = -generalized_box_iou(out_bbox, tgt_bbox) # Final cost matrix C = self.cost_bbox * cost_bbox + self.cost_class * cost_class + self.cost_giou * cost_giou C = C.view(bs, num_queries, -1).cpu() sizes = [len(v["boxes"]) for v in targets] indices = [ linear_sum_assignment(c[i]) for i, c in enumerate(C.split(sizes, -1)) ] return [(torch.as_tensor(i, dtype=torch.int64), torch.as_tensor(j, dtype=torch.int64)) for i, j in indices]
# transformed = transform_to_tensor(pil_img) # pil_img = ImageOps.grayscale(pil_img) transformed = transform_to_tensor(pil_img) images.append(transformed) images = torch.stack(images) positions = torch.Tensor(positions) torch.save(images, os.path.join(root_path_save,'{}.pt'.format('images'))) torch.save(positions, os.path.join(root_path_save,'{}.pt'.format('deg_pos'))) distances = torch.cdist(images.view(images.shape[0],3*480*320),images.view(images.shape[0],3*480*320)) print(distances.shape) full_dataset = TensorDataset(images,positions) # # #ind = random.sample(range(180),180) # ind_train = random.sample(range(180),120) # ind_test = random.sample(range(180),120) # print(positions[ind_train].unsqueeze_(0).shape) # # train_dataset = TensorDataset(images[ind_train,:,:,:],positions[ind_train]) # test_dataset = TensorDataset(images[ind_test,:,:,:],positions[ind_test]) torch.save(full_dataset, os.path.join(root_path_save,'{}.pt'.format('full_dataset'))) # torch.save(train_dataset, os.path.join(root_path_save,'{}.pt'.format('train_dataset'))) # torch.save(test_dataset, os.path.join(root_path_save,'{}.pt'.format('test_dataset'))) # dataloader = DataLoader(full_dataset, batch_size=images.shape[0], pin_memory=True, drop_last=True,shuffle=False)
def _kernel(x, y, basis_func, sigma): return basis_func(torch.cdist(x, y + EPSILON) * torch.abs(sigma))
def match_smnn( desc1: torch.Tensor, desc2: torch.Tensor, th: float = 0.8, dm: Optional[torch.Tensor] = None ) -> Tuple[torch.Tensor, torch.Tensor]: """Function, which finds mutual nearest neighbors in desc2 for each vector in desc1. the method satisfies first to second nearest neighbor distance <= th. If the distance matrix dm is not provided, :py:func:`torch.cdist` is used. Args: desc1: Batch of descriptors of a shape :math:`(B1, D)`. desc2: Batch of descriptors of a shape :math:`(B2, D)`. th: distance ratio threshold. dm: Tensor containing the distances from each descriptor in desc1 to each descriptor in desc2, shape of :math:`(B1, B2)`. Return: - Descriptor distance of matching descriptors, shape of. :math:`(B3, 1)`. - Long tensor indexes of matching descriptors in desc1 and desc2, shape of :math:`(B3, 2)` where 0 <= B3 <= B1. """ if len(desc1.shape) != 2: raise AssertionError if len(desc2.shape) != 2: raise AssertionError if desc1.shape[0] < 2: raise AssertionError if desc2.shape[0] < 2: raise AssertionError if dm is None: dm = torch.cdist(desc1, desc2) else: if not ((dm.size(0) == desc1.size(0)) and (dm.size(1) == desc2.size(0))): raise AssertionError dists1, idx1 = match_snn(desc1, desc2, th, dm) dists2, idx2 = match_snn(desc2, desc1, th, dm.t()) if len(dists2) > 0 and len(dists1) > 0: idx2 = idx2.flip(1) idxs_dm = torch.cdist(idx1.float(), idx2.float(), p=1) mutual_idxs1 = idxs_dm.min(dim=1)[0] < 1e-8 mutual_idxs2 = idxs_dm.min(dim=0)[0] < 1e-8 good_idxs1 = idx1[mutual_idxs1.view(-1)] good_idxs2 = idx2[mutual_idxs2.view(-1)] dists1_good = dists1[mutual_idxs1.view(-1)] dists2_good = dists2[mutual_idxs2.view(-1)] _, idx_upl1 = torch.sort(good_idxs1[:, 0]) _, idx_upl2 = torch.sort(good_idxs2[:, 0]) good_idxs1 = good_idxs1[idx_upl1] match_dists = torch.max(dists1_good[idx_upl1], dists2_good[idx_upl2]) matches_idxs = good_idxs1 else: matches_idxs, match_dists = torch.empty( 0, 2, device=dm.device), torch.empty(0, 1, device=dm.device) return match_dists.view(-1, 1), matches_idxs.view(-1, 2)
def train(self): train_set = CIFAR10(self.config.data_dir, train=True, transform=general_transform['train'], download=True) loader = DataLoader(train_set, self.config.batch_size, shuffle=True, num_workers=self.config.num_workers) optimizer = optim.SGD(self.model.parameters(), lr=self.config.lr, momentum=self.config.momentum, weight_decay=self.config.weight_decay) criterion = nn.CrossEntropyLoss() classes = torch.Tensor([i for i in range(10)]).float().to(self.device) for epoch in range(1, self.config.epoch + 1): log_dict = { 'batch_loss': .0, 'acc': .0, 'lr': .0 } min_loss = 100 max_acc = 0 total = 0 correct = 0 for inputs, targets in tqdm(loader): inputs, targets = self.to_device(inputs, targets) output, _ = self.model(inputs) # output : batch x 512 x 1 x 1 # calculate l2 distance and use squared distance. dist = torch.cdist(self.model.centers, output.squeeze()) ** 2 dist = dist.T # output shape must be Batch x Class r_target = targets.repeat(self.config.num_classes).view(self.config.num_classes, -1) mask = (r_target == classes.unsqueeze(1)).T.float() print(targets.shape) print(r_target.shape) print(mask.shape) exit() inclass_loss = dist * mask # optimizer.zero_grad() # loss.backward() # optimizer.step() # # # _, pred = center_sim.T.max(1) # # total += targets.size(0) # correct += pred.eq(targets).sum() # log_dict['batch_loss'] += loss.item() log_dict['acc'] = correct / total log_dict['batch_loss'] /= len(loader) print_log(f"{epoch} / {self.config.epoch}", log_dict)
def L2_dist(c, x): return torch.cdist(c, x) ** 2
included_labels['display_name'][i]: pd.Series(index=included_labels['display_name'][bottom[:, i]].values, data=leastlike[i][bottom[:, i]]) for i in range(leastlike.shape[0]) }) del leastlike del cc gc.collect() else: assert metric == 'distance' wt = emb.wi_mu.weight.detach() wt = wt.to(device) # compute_mode="donot_use_mm_for_euclid_dist" is required or else # the distance between something and itself is not 0 # don't know why. dist = torch.cdist(wt, wt, compute_mode="donot_use_mm_for_euclid_dist") dist = dist.detach().cpu() # same as above, replace values of 0 with inf or -inf so we can sort mostlike = dist.clone() mostlike[torch.isclose(dist, torch.tensor([0.0]))] = np.inf mostlike = mostlike.numpy() top = mostlike.argsort(axis=0)[:5].copy() # make into data structures maxes = pd.Series({ included_labels['display_name'][i]: pd.Series(index=included_labels['display_name'][top[:, i]].values, data=mostlike[i][top[:, i]]) for i in range(wt.shape[0]) }) del mostlike
def hamming(a, b): return torch.cdist(a, b, p=0)
# print(input) flatten = input.permute([0, 2, 3, 1]).contiguous().view( -1, self.embed_dim).type(torch.float) dist = torch.cdist(flatten, self.embed.weight.data) _, embed_idx = dist.min(dim=1) # print(embed_idx) embed_idx = embed_idx.view(N, H, W) quantize = self.embed_code(embed_idx) quantize = quantize.permute([0, 3, 1, 2]) # print(quantize) embed_loss = F.mse_loss(input.detach(), quantize) commit_loss = F.mse_loss(quantize.detach(), input) quantize = input + (quantize - input).detach() return quantize, embed_idx, embed_loss, commit_loss # unit test if __name__ == '__main__': codebook = nn.Embedding(num_embeddings=3, embedding_dim=2) codebook.weight.data.uniform_(0, 10) test_vec = torch.randint(high=10, size=(1, 2)).type(torch.float) print(codebook.weight.data) print(test_vec) dist = torch.cdist(test_vec, codebook.weight.data) print(dist) _, idx = dist.min(dim=1) print(idx) print(codebook(idx))
def torch_cdist(X): tensor = torch.from_numpy(X).to(torch.device('cuda')) d = torch.cdist(tensor, tensor, p=2) d = d.cpu().numpy() return d
def compute_l2_dist(x1, x2): """compute an (m, n) L2 distance matrix from (m, d) and (n, d) matrices""" return torch.cdist(x1.unsqueeze(0), x2.unsqueeze(0), p=2).squeeze(0).pow(2)
def forward(self, outputs, targets): """ Performs the matching Params: outputs: This is a dict that contains at least these entries: "pred_logits": Tensor of dim [batch_size, num_queries, num_classes] with the classification logits "pred_boxes": Tensor of dim [batch_size, num_queries, 4] with the predicted box coordinates targets: This is a list of targets (len(targets) = batch_size), where each target is a dict containing: "labels": Tensor of dim [num_target_boxes] (where num_target_boxes is the number of ground-truth objects in the target) containing the class labels "boxes": Tensor of dim [num_target_boxes, 4] containing the target box coordinates Returns: A list of size batch_size, containing tuples of (index_i, index_j) where: - index_i is the indices of the selected predictions (in order) - index_j is the indices of the corresponding selected targets (in order) For each batch element, it holds: len(index_i) = len(index_j) = min(num_queries, num_target_boxes) """ bs, k, h, w = outputs["pred_logits"].shape # We flatten to compute the cost matrices in a batch batch_out_prob = outputs["pred_logits"].permute(0, 2, 3, 1).reshape( bs, h * w, k).sigmoid() # [batch_size, num_queries, num_classes] batch_out_bbox = outputs["pred_boxes"].permute(0, 2, 3, 1).reshape( bs, h * w, 4) # [batch_size, num_queries, 4] indices = [] for i in range(bs): tgt_ids = targets[i]["labels"] if tgt_ids.shape[0] == 0: indices.append(([], [])) continue tgt_bbox = targets[i]["boxes_xyxy"] out_prob = batch_out_prob[i] out_bbox = batch_out_bbox[i] # Compute the classification cost. alpha = self.focal_loss_alpha gamma = self.focal_loss_gamma neg_cost_class = (1 - alpha) * (out_prob**gamma) * ( -(1 - out_prob + 1e-8).log()) pos_cost_class = alpha * ( (1 - out_prob)**gamma) * (-(out_prob + 1e-8).log()) cost_class = pos_cost_class[:, tgt_ids] - neg_cost_class[:, tgt_ids] # Compute the L1 cost between boxes image_size_out = targets[i]["image_size_xyxy"].unsqueeze(0).repeat( h * w, 1) image_size_tgt = targets[i]["image_size_xyxy_tgt"] out_bbox_ = out_bbox / image_size_out tgt_bbox_ = tgt_bbox / image_size_tgt cost_bbox = torch.cdist(out_bbox_, tgt_bbox_, p=1) # Compute the giou cost betwen boxes cost_giou = -generalized_box_iou(out_bbox, tgt_bbox) # Final cost matrix C = self.cost_bbox * cost_bbox + self.cost_class * cost_class + self.cost_giou * cost_giou _, src_ind = torch.min(C, dim=0) tgt_ind = torch.arange(len(tgt_ids)).to(src_ind) indices.append((src_ind, tgt_ind)) return [(torch.as_tensor(i, dtype=torch.int64), torch.as_tensor(j, dtype=torch.int64)) for i, j in indices]
def forward(self, input): return -torch.cdist(input, self.weight, p=self.p)
def train(model, epoch): since = time.time() # AVERAGE METER losses = AverageMeter() # TRIGGER PARAMETERS trans_image = transforms.Compose([transforms.Resize((224, 224)), transforms.ToTensor(), ]) trans_trigger = transforms.Compose([transforms.Resize((patch_size, patch_size)), transforms.ToTensor(), ]) # PERTURBATION PARAMETERS eps1 = (eps/255.0) lr1 = lr trigger = Image.open('data/triggers/trigger_{}.png'.format(trigger_id)).convert('RGB') trigger = trans_trigger(trigger).unsqueeze(0)#.cuda(gpu) # SOURCE AND TARGET DATASETS target_filelist = "ImageNet_data_list/poison_generation/" + target_wnid + ".txt" # Use source wnid list if num_source==1: logging.info("Using single source for this experiment.") else: logging.info("Using multiple source for this experiment.") with open("data/{}/multi_source_filelist.txt".format(experimentID),"w") as f1: with open(source_wnid_list) as f2: source_wnids = f2.readlines() source_wnids = [s.strip() for s in source_wnids] for source_wnid in source_wnids: with open("ImageNet_data_list/poison_generation/" + source_wnid + ".txt", "r") as f2: shutil.copyfileobj(f2, f1) source_filelist = "data/{}/multi_source_filelist.txt".format(experimentID) dataset_target = PoisonGenerationDataset(data_root + "/train", target_filelist, trans_image) dataset_source = PoisonGenerationDataset(data_root + "/train", source_filelist, trans_image) # SOURCE AND TARGET DATALOADERS train_loader_target = torch.utils.data.DataLoader(dataset_target, batch_size=100, shuffle=True, num_workers=8, pin_memory=True) train_loader_source = torch.utils.data.DataLoader(dataset_source, batch_size=100, shuffle=True, num_workers=8, pin_memory=True) logging.info("Number of target images:{}".format(len(dataset_target))) logging.info("Number of source images:{}".format(len(dataset_source))) # USE ITERATORS ON DATALOADERS TO HAVE DISTINCT PAIRING EACH TIME iter_target = iter(train_loader_target) iter_source = iter(train_loader_source) num_poisoned = 0 for i in range(len(train_loader_target)): # LOAD ONE BATCH OF SOURCE AND ONE BATCH OF TARGET (input1, path1) = next(iter_source) (input2, path2) = next(iter_target) img_ctr = 0 input1 = input1#.cuda(gpu) input2 = input2#.cuda(gpu) pert = nn.Parameter(torch.zeros_like(input2, requires_grad=True))#.cuda(gpu)) for z in range(input1.size(0)): if not rand_loc: start_x = 224-patch_size-5 start_y = 224-patch_size-5 else: start_x = random.randint(0, 224-patch_size-1) start_y = random.randint(0, 224-patch_size-1) # PASTE TRIGGER ON SOURCE IMAGES input1[z, :, start_y:start_y+patch_size, start_x:start_x+patch_size] = trigger output1, feat1 = model(input1) feat1 = feat1.detach().clone() for k in range(input1.size(0)): img_ctr = img_ctr+1 # input2_pert = (pert[k].clone().cpu()) fname = saveDir_patched + '/' + 'badnet_' + str(os.path.basename(path1[k])).split('.')[0] + '_' + 'epoch_' + str(epoch).zfill(2)\ + str(img_ctr).zfill(5)+'.png' save_image(input1[k].clone().cpu(), fname) num_poisoned +=1 for j in range(num_iter): lr1 = adjust_learning_rate(lr, j) output2, feat2 = model(input2+pert) # FIND CLOSEST PAIR WITHOUT REPLACEMENT feat11 = feat1.clone() dist = torch.cdist(feat1, feat2) for _ in range(feat2.size(0)): dist_min_index = (dist == torch.min(dist)).nonzero().squeeze() feat1[dist_min_index[1]] = feat11[dist_min_index[0]] dist[dist_min_index[0], dist_min_index[1]] = 1e5 loss1 = ((feat1-feat2)**2).sum(dim=1) loss = loss1.sum() losses.update(loss.item(), input1.size(0)) loss.backward() pert = pert- lr1*pert.grad pert = torch.clamp(pert, -eps1, eps1).detach_() pert = pert + input2 pert = pert.clamp(0, 1) if j%10 == 0: logging.info("Epoch: {:2d} | i: {} | iter: {:5d} | LR: {:2.4f} | Loss Val: {:5.3f} | Loss Avg: {:5.3f}" .format(epoch, i, j, lr1, losses.val, losses.avg)) if loss1.max().item() < 10 or j == (num_iter-1): for k in range(input2.size(0)): img_ctr = img_ctr+1 input2_pert = (pert[k].clone().cpu()) fname = saveDir_poison + '/' + 'loss_' + str(int(loss1[k].item())).zfill(5) + '_' + 'epoch_' + \ str(epoch).zfill(2) + '_' + str(os.path.basename(path2[k])).split('.')[0] + '_' + \ str(os.path.basename(path1[k])).split('.')[0] + '_kk_' + str(img_ctr).zfill(5)+'.png' save_image(input2_pert, fname) num_poisoned +=1 break pert = pert - input2 pert.requires_grad = True time_elapsed = time.time() - since logging.info('Training complete one epoch in {:.0f}m {:.0f}s'.format(time_elapsed // 60, time_elapsed % 60))
def forward(self, outputs, targets): """ Performs the matching Params: outputs: This is a dict that contains at least these entries: "pred_logits": Tensor of dim [batch_size, num_queries, num_classes] with the classification logits "pred_boxes": Tensor of dim [batch_size, num_queries, 4] with the predicted box coordinates targets: This is a list of targets (len(targets) = batch_size), where each target is a dict containing: "labels": Tensor of dim [num_target_boxes] (where num_target_boxes is the number of ground-truth objects in the target) containing the class labels "boxes": Tensor of dim [num_target_boxes, 4] containing the target box coordinates Returns: A list of size batch_size, containing tuples of (index_i, index_j) where: - index_i is the indices of the selected predictions (in order) - index_j is the indices of the corresponding selected targets (in order) For each batch element, it holds: len(index_i) = len(index_j) = min(num_queries, num_target_boxes) """ with torch.no_grad(): bs, num_queries = outputs["pred_logits"].shape[:2] # We flatten to compute the cost matrices in a batch out_prob = outputs["pred_logits"].flatten(0, 1).sigmoid() out_bbox = outputs["pred_boxes"].flatten( 0, 1) # [batch_size * num_queries, 4] # Also concat the target labels and boxes tgt_ids = torch.cat([v["labels"] for v in targets]) tgt_bbox = torch.cat([v["boxes"] for v in targets]) # Compute the classification cost. alpha = 0.25 gamma = 2.0 neg_cost_class = (1 - alpha) * (out_prob**gamma) * ( -(1 - out_prob + 1e-8).log()) pos_cost_class = alpha * ( (1 - out_prob)**gamma) * (-(out_prob + 1e-8).log()) cost_class = pos_cost_class[:, tgt_ids] - neg_cost_class[:, tgt_ids] # Compute the L1 cost between boxes cost_bbox = torch.cdist(out_bbox, tgt_bbox, p=1) # Compute the giou cost betwen boxes cost_giou = -generalized_box_iou(box_cxcywh_to_xyxy(out_bbox), box_cxcywh_to_xyxy(tgt_bbox)) # Final cost matrix C = self.cost_bbox * cost_bbox + self.cost_class * cost_class + self.cost_giou * cost_giou C = C.view(bs, num_queries, -1).cpu() sizes = [len(v["boxes"]) for v in targets] indices = [ linear_sum_assignment(c[i]) for i, c in enumerate(C.split(sizes, -1)) ] return [(torch.as_tensor(i, dtype=torch.int64), torch.as_tensor(j, dtype=torch.int64)) for i, j in indices]
def assign_anchors(coords, coords_ref, dist_thr=None, return_perm=False, cdist=None): """ Assign the closest anchors with coords coords_ref - cdist: precomputed distance matrix between coords and coords_ref # Test with equal shape for coords and coords_ref >>> coords_ref = torch.tensor([[0., 0., 0.], [1., 2., 3.], [4., 5., 6.], [7., 8., 9.]]) >>> coords = torch.zeros_like(coords_ref) >>> coords[0] = coords_ref[1] >>> coords[1] = coords_ref[3] >>> coords[2] = coords_ref[0] >>> coords[3] = coords_ref[2] >>> coords tensor([[1., 2., 3.], [7., 8., 9.], [0., 0., 0.], [4., 5., 6.]]) >>> get_RMSD(coords, coords_ref) tensor(7.5166) >>> assignment, sel, P = assign_anchors(coords, coords_ref, return_perm=True) >>> coords_ordered = coords[assignment] >>> (coords.T.mm(P).T == coords_ordered).all() tensor(True) >>> coords_ordered tensor([[0., 0., 0.], [1., 2., 3.], [4., 5., 6.], [7., 8., 9.]]) >>> get_RMSD(coords_ordered, coords_ref[sel]) tensor(0.) # Test with size of coords_ref lower than coords >>> coords_ref = torch.tensor([[-1., -2., -3.], [1., 2., 3.], [4., 5., 6.], [7., 8., 9.]]) >>> coords = torch.zeros((5, 3)) >>> coords[0] = coords_ref[1] >>> coords[1] = coords_ref[3] >>> coords[4] = coords_ref[0] >>> coords[3] = coords_ref[2] >>> coords[2] = torch.tensor([10., 11., 12.]) >>> coords tensor([[ 1., 2., 3.], [ 7., 8., 9.], [10., 11., 12.], [ 4., 5., 6.], [-1., -2., -3.]]) >>> assignment, sel, P = assign_anchors(coords, coords_ref, return_perm=True) >>> coords_ordered = coords[assignment] >>> (coords.T.mm(P).T == coords_ordered).all() tensor(True) >>> coords_ordered tensor([[-1., -2., -3.], [ 1., 2., 3.], [ 4., 5., 6.], [ 7., 8., 9.]]) >>> get_RMSD(coords_ordered, coords_ref[sel]) tensor(0.) # Test with size of coords_ref higher than coords >>> coords_ref = torch.tensor([[-1., -2., -3.], [1., 2., 3.], [4., 5., 6.], [7., 8., 9.], [10, 11, 12]]) >>> coords = torch.zeros((4, 3)) >>> coords[0] = coords_ref[1] >>> coords[1] = coords_ref[3] >>> coords[2] = coords_ref[0] >>> coords[3] = coords_ref[4] >>> coords tensor([[ 1., 2., 3.], [ 7., 8., 9.], [-1., -2., -3.], [10., 11., 12.]]) >>> assignment, sel, P = assign_anchors(coords, coords_ref, return_perm=True) >>> coords_ordered = coords[assignment] >>> (coords.T.mm(P).T == coords_ordered).all() tensor(True) >>> coords_ordered tensor([[-1., -2., -3.], [ 1., 2., 3.], [ 7., 8., 9.], [10., 11., 12.]]) >>> sel array([0, 1, 3, 4]) >>> assignment array([2, 0, 1, 3]) >>> get_RMSD(coords_ordered, coords_ref[sel]) tensor(0.) >>> coords[2] += 100. >>> assignment, sel, P = assign_anchors(coords, coords_ref, dist_thr=4., return_perm=True) >>> coords_ref tensor([[-1., -2., -3.], [ 1., 2., 3.], [ 4., 5., 6.], [ 7., 8., 9.], [10., 11., 12.]]) >>> coords tensor([[ 1., 2., 3.], [ 7., 8., 9.], [99., 98., 97.], [10., 11., 12.]]) >>> coords_ordered = coords[assignment] >>> coords_ordered tensor([[ 1., 2., 3.], [ 7., 8., 9.], [10., 11., 12.]]) >>> get_RMSD(coords_ordered, coords_ref[sel]) tensor(0.) >>> coords.T.mm(P).T tensor([[ 1., 2., 3.], [ 7., 8., 9.], [10., 11., 12.]]) """ if cdist is None: cdist = torch.cdist(coords, coords_ref) cdist = cdist.cpu().numpy() if dist_thr is not None: cdist[cdist > dist_thr] = 9999.99 row_ind, col_ind = scipy.optimize.linear_sum_assignment(cdist) if dist_thr is not None: distances = cdist[row_ind, col_ind] # sel = distances <= dist_thr sel = distances < 9999.99 row_ind = row_ind[sel] col_ind = col_ind[sel] n = coords.shape[0] n_ref = coords_ref.shape[0] assignment = -np.ones(n_ref, dtype=int) assignment[col_ind] = row_ind assignment = assignment[assignment > -1] sel = list(col_ind) sel.sort() sel = np.asarray(sel) if return_perm: P = torch.zeros((n, n_ref)) P[assignment, torch.arange(len(assignment))] = 1. P = P[:, :n] # P[:, P.sum(axis=0) == 0] = 1 / n P = P[:, P.sum(axis=0) != 0] return assignment, sel, P else: return assignment, sel
def re_ranking_old(probFea, galFea, k1, k2, lambda_value, local_distmat=None, theta_value=0.5, only_local=False, MemorySave=True, Minibatch=5000): assert k2 <= k1 + 1 # if feature vector is numpy, you should use 'torch.tensor' transform it to tensor query_num = probFea.size(0) all_num = query_num + galFea.size(0) if only_local: original_dist = local_distmat else: print('Computing original distance using GPU ...') feat = torch.cat([probFea, galFea]).cuda() if MemorySave: distmat = torch.zeros( (all_num, all_num), dtype=torch.float16) # 14 GB memory on Round2 i = 0 while True: it = i + Minibatch #print('i, it', i, it) if it < feat.size()[0]: distmat[i:it, :] = torch.pow( torch.cdist(feat[i:it, :], feat), 2) else: distmat[i:, :] = torch.pow(torch.cdist(feat[i:, :], feat), 2) break i = it else: ### new API distmat = torch.pow(torch.cdist(feat, feat), 2) #print('Copy distmat to original_dist ...') original_dist = distmat.numpy() # 14 GB memory del distmat del feat if not local_distmat is None: original_dist = original_dist * theta_value + local_distmat * ( 1 - theta_value) gallery_num = original_dist.shape[0] ### memory optimization print('Division ...') # memory inefficient #original_dist = np.transpose(original_dist / np.max(original_dist, axis=0)) m = np.max(original_dist, axis=0) original_dist = mem_saving_divide(original_dist, m) #print('Transpose ...') original_dist = original_dist.T # ultra fast ### print('Argsort to get initial_rank ...') #initial_rank = np.argsort(original_dist) # .astype(np.int32) initial_rank = mem_saving_argsort(original_dist, top_k=k1 + 1) # Time: 8.5 min print('Allocate V ...') V = np.zeros_like(original_dist, dtype=np.float16) # 14 GB memory print('Start re_ranking ...') for i in tqdm.tqdm(range(all_num)): # Time: 18 s # k-reciprocal neighbors forward_k_neigh_index = initial_rank[i, :k1 + 1] backward_k_neigh_index = initial_rank[forward_k_neigh_index, :k1 + 1] fi = np.where(backward_k_neigh_index == i)[0] k_reciprocal_index = forward_k_neigh_index[fi] k_reciprocal_expansion_index = k_reciprocal_index for j in range(len(k_reciprocal_index)): candidate = k_reciprocal_index[j] candidate_forward_k_neigh_index = initial_rank[ candidate, :int(np.around(k1 / 2)) + 1] candidate_backward_k_neigh_index = initial_rank[ candidate_forward_k_neigh_index, :int(np.around(k1 / 2)) + 1] fi_candidate = np.where( candidate_backward_k_neigh_index == candidate)[0] candidate_k_reciprocal_index = candidate_forward_k_neigh_index[ fi_candidate] if len( np.intersect1d(candidate_k_reciprocal_index, k_reciprocal_index) ) > 2 / 3 * len(candidate_k_reciprocal_index): k_reciprocal_expansion_index = np.append( k_reciprocal_expansion_index, candidate_k_reciprocal_index) k_reciprocal_expansion_index = np.unique(k_reciprocal_expansion_index) weight = np.exp(-original_dist[i, k_reciprocal_expansion_index]) V[i, k_reciprocal_expansion_index] = weight / np.sum(weight) original_dist = original_dist[:query_num, :] if k2 != 1: V_qe = np.zeros_like(V, dtype=np.float16) # 14 GB memory for i in tqdm.tqdm(range(all_num)): # Time: 12.6 min V_qe[i, :] = np.mean(V[initial_rank[i, :k2], :], axis=0) V = V_qe del V_qe del initial_rank invIndex = [] for i in tqdm.tqdm(range(gallery_num)): # Time: 20 s invIndex.append(np.where(V[:, i] != 0)[0]) ##### To save memory, don't allocate another matrix, re-use memory of original_dist instead #jaccard_dist = np.zeros_like(original_dist, dtype=np.float16) for i in tqdm.tqdm(range(query_num)): # Time: 1 min temp_min = np.zeros(shape=[1, gallery_num], dtype=np.float16) indNonZero = np.where(V[i, :] != 0)[0] indImages = [invIndex[ind] for ind in indNonZero] for j in range(len(indNonZero)): temp_min[0, indImages[j]] = temp_min[0, indImages[j]] + np.minimum( V[i, indNonZero[j]], V[indImages[j], indNonZero[j]]) ####### #jaccard_dist[i] = 1 - temp_min / (2 - temp_min) jaccard_dist_i = 1 - temp_min / (2 - temp_min) original_dist[i] = jaccard_dist_i * ( 1 - lambda_value) + original_dist[i] * lambda_value ####### #final_dist = jaccard_dist * (1 - lambda_value) + original_dist * lambda_value #del original_dist del V #del jaccard_dist final_dist = original_dist[:query_num, query_num:] return final_dist
def get_k_hamming_neighbours(enc_train, enc_query_imgs): hammingDistances = torch.cdist(enc_query_imgs, enc_train, p=0) sortedDistances, indices = torch.sort(hammingDistances) return indices
def tool_batch(args): for cval in range(4): if args.resume.format( cval) == args.resume: # no cross validation is used # skip all other datasets if cval != args.cross_val: continue args.resume = args.resume.format(cval) model = SpineDETR(args) df = pd.read_csv(f'{args.spine_ann_2d}/csv_test_{cval}.csv') df = df[['patient_id', 'cross_val', 'filename']].drop_duplicates() for row_i, (index, row) in enumerate(df.iterrows()): img_path = args.spine_imgs_2d + f"{row['patient_id']}/{row['filename']}" all_logits, all_centers = [], [] for batch, windows, src_img in batch_gen( img_path, stride=args.stride, max_batch_size=args.batch_size, resize=args.resize): out = model(batch) all_logits.append(out['pred_logits']) all_centers.append(out['pred_boxes']) print( f' CVAL: {cval} progress: {row_i}/{len(df)} id: {row["patient_id"]} windows# {len(all_logits)}{" " * 10}', end='\r') all_logits = torch.cat(all_logits) all_centers = torch.cat(all_centers) window_size = args.rand_crop window_centers = torch.Tensor(windows).to(args.device) window_centers += window_size // 2 out_centers = [] for i, (window, logits, centers) in enumerate(zip(windows, all_logits, all_centers)): centers = centers[logits.squeeze() > 0.5] centers = centers * window_size centers[:, 0] += window[0] centers[:, 1] += window[1] # out_centers.append(centers) if centers.numel() > 0: dist = torch.cdist(centers, window_centers) min_idx = dist.argmin(1) correct_centers = centers[min_idx == i] out_centers.append(correct_centers) out_centers = torch.cat(out_centers) os.makedirs(f'{args.output_dir}/{row["patient_id"]}', exist_ok=True) sorted_idx = torch.argsort(out_centers[:, 1]) out_centers = out_centers[sorted_idx] out_dict = {} for i, (x, y) in enumerate(out_centers): out_dict[str(i)] = [{'pos': [x.item(), y.item()]}] with open( f'{args.output_dir}/{row["patient_id"]}/{row["patient_id"]}.json', 'w') as f: json.dump(out_dict, f) width, height = F._get_image_size(src_img) out_centers[:, 0] /= width out_centers[:, 1] /= height out_image = spine_plot_centers(src_img, out_centers) out_image.save( f'{args.output_dir}/{row["patient_id"]}/{row["patient_id"]}.jpg' ) print( f' CVAL: {cval} progress: {row_i}/{len(df)} id: {row["patient_id"]} windows# {len(all_logits)} MEM Res: {torch.cuda.memory_reserved() / (2 ** 30):.02f} Gb{" " * 10}', end='\n')
def dist_mat(X, squared=True): D = torch.cdist(X, X) return D if not squared else D.sqrt()
def model(self, data): ''' Define the parameters ''' n_ind = data['n_ind'] n_trt = data['n_trt'] n_tms = data['n_tms'] n_mrk = data['n_mrk'] n_prs = n_trt * n_tms * n_mrk plt_ind = pyro.plate('individuals', n_ind, dim=-3) plt_trt = pyro.plate('treatments', n_trt, dim=-2) plt_tms = pyro.plate('times', n_tms, dim=-1) pars = {} # covariance factors with plt_tms: # learning dt time step sizes # if k(t1,t2) is independent of time, can instead learn scales and variances for RBF kernels that use data['time_vals'] pars['dt0'] = pyro.sample('dt0', dist.Normal(0, 1)) pars['dt1'] = pyro.sample('dt1', dist.Normal(0, 1)) pars['theta_trt0'] = pyro.sample('theta_trt0', dist.HalfCauchy(torch.ones(n_trt))) pars['theta_mrk0'] = pyro.sample('theta_mrk0', dist.HalfCauchy(torch.ones(n_mrk))) pars['theta_trt1'] = pyro.sample('theta_trt1', dist.HalfCauchy(torch.ones(n_trt))) pars['L_omega_trt1'] = pyro.sample( 'L_omega_trt1', dist.LKJCorrCholesky(n_trt, torch.ones(1))) pars['theta_mrk1'] = pyro.sample('theta_mrk1', dist.HalfCauchy(torch.ones(n_mrk))) pars['L_omega_mrk1'] = pyro.sample( 'L_omega_mrk1', dist.LKJCorrCholesky(n_mrk, torch.ones(1))) times0 = fun.pad(torch.cumsum(pars['dt0'].exp().log1p(), 0), (1, 0), value=0)[:-1].unsqueeze(1) times1 = fun.pad(torch.cumsum(pars['dt1'].exp().log1p(), 0), (1, 0), value=0)[:-1].unsqueeze(1) cov_t0 = (-torch.cdist(times0, times0)).exp() cov_t1 = (-torch.cdist(times1, times1)).exp() cov_i0 = pars['theta_trt0'].diag() L_Omega_trt = torch.mm(torch.diag(pars['theta_trt1'].sqrt()), pars['L_omega_trt1']) cov_i1 = L_Omega_trt.mm(L_Omega_trt.t()) cov_m0 = pars['theta_mrk0'].diag() L_Omega_mrk = torch.mm(torch.diag(pars['theta_mrk1'].sqrt()), pars['L_omega_mrk1']) cov_m1 = L_Omega_mrk.mm(L_Omega_mrk.t()) # kronecker product of the factors cov_itm0 = torch.einsum('ij,tu,mn->itmjun', [cov_i0, cov_t0, cov_m0]).view(n_prs, n_prs) cov_itm1 = torch.einsum('ij,tu,mn->itmjun', [cov_i1, cov_t1, cov_m1]).view(n_prs, n_prs) # global and individual level params of each marker, treatment, and time point pars['glb'] = pyro.sample( 'glb', dist.MultivariateNormal(torch.zeros(n_prs), cov_itm0)) with plt_ind: pars['ind'] = pyro.sample( 'ind', dist.MultivariateNormal(torch.zeros(n_prs), cov_itm1)) # observation noise, time series bias and scale pars['noise_scale'] = pyro.sample('noise_scale', dist.HalfCauchy(torch.ones(n_mrk))) pars['t0_scale'] = pyro.sample('t0_scale', dist.HalfCauchy(torch.ones(n_mrk))) with plt_ind: pars['t0'] = pyro.sample( 't0', dist.MultivariateNormal(torch.zeros(n_mrk), pars['t0_scale'].diag())) with plt_trt, plt_tms: pars['noise'] = pyro.sample( 'noise', dist.MultivariateNormal(torch.zeros(n_mrk), pars['noise_scale'].diag())) # likelihood of the data distr = self.get_distr(data, pars) pyro.sample('obs', distr, obs=data['Y'])
def assign(self, bbox_pred, anchor, gt_bboxes, gt_bboxes_ignore=None, gt_labels=None): num_gts, num_bboxes = gt_bboxes.size(0), bbox_pred.size(0) # 1. assign -1 by default assigned_gt_inds = bbox_pred.new_full((num_bboxes, ), 0, dtype=torch.long) assigned_labels = bbox_pred.new_full((num_bboxes, ), -1, dtype=torch.long) if num_gts == 0 or num_bboxes == 0: # No ground truth or boxes, return empty assignment if num_gts == 0: # No ground truth, assign all to background assigned_gt_inds[:] = 0 assign_result = AssignResult( num_gts, assigned_gt_inds, None, labels=assigned_labels) assign_result.set_extra_property( 'pos_idx', bbox_pred.new_empty(0, dtype=torch.bool)) assign_result.set_extra_property('pos_predicted_boxes', bbox_pred.new_empty((0, 4))) assign_result.set_extra_property('target_boxes', bbox_pred.new_empty((0, 4))) return assign_result # 2. Compute the L1 cost between boxes # Note that we use anchors and predict boxes both cost_bbox = torch.cdist( bbox_xyxy_to_cxcywh(bbox_pred), bbox_xyxy_to_cxcywh(gt_bboxes), p=1) cost_bbox_anchors = torch.cdist( bbox_xyxy_to_cxcywh(anchor), bbox_xyxy_to_cxcywh(gt_bboxes), p=1) # We found that topk function has different results in cpu and # cuda mode. In order to ensure consistency with the source code, # we also use cpu mode. # TODO: Check whether the performance of cpu and cuda are the same. C = cost_bbox.cpu() C1 = cost_bbox_anchors.cpu() # self.match_times x n index = torch.topk( C, # c=b,n,x c[i]=n,x k=self.match_times, dim=0, largest=False)[1] # self.match_times x n index1 = torch.topk(C1, k=self.match_times, dim=0, largest=False)[1] # (self.match_times*2) x n indexes = torch.cat((index, index1), dim=1).reshape(-1).to(bbox_pred.device) pred_overlaps = self.iou_calculator(bbox_pred, gt_bboxes) anchor_overlaps = self.iou_calculator(anchor, gt_bboxes) pred_max_overlaps, _ = pred_overlaps.max(dim=1) anchor_max_overlaps, _ = anchor_overlaps.max(dim=0) # 3. Compute the ignore indexes use gt_bboxes and predict boxes ignore_idx = pred_max_overlaps > self.neg_ignore_thr assigned_gt_inds[ignore_idx] = -1 # 4. Compute the ignore indexes of positive sample use anchors # and predict boxes pos_gt_index = torch.arange( 0, C1.size(1), device=bbox_pred.device).repeat(self.match_times * 2) pos_ious = anchor_overlaps[indexes, pos_gt_index] pos_ignore_idx = pos_ious < self.pos_ignore_thr pos_gt_index_with_ignore = pos_gt_index + 1 pos_gt_index_with_ignore[pos_ignore_idx] = -1 assigned_gt_inds[indexes] = pos_gt_index_with_ignore if gt_labels is not None: assigned_labels = assigned_gt_inds.new_full((num_bboxes, ), -1) pos_inds = torch.nonzero( assigned_gt_inds > 0, as_tuple=False).squeeze() if pos_inds.numel() > 0: assigned_labels[pos_inds] = gt_labels[ assigned_gt_inds[pos_inds] - 1] else: assigned_labels = None assign_result = AssignResult( num_gts, assigned_gt_inds, anchor_max_overlaps, labels=assigned_labels) assign_result.set_extra_property('pos_idx', ~pos_ignore_idx) assign_result.set_extra_property('pos_predicted_boxes', bbox_pred[indexes]) assign_result.set_extra_property('target_boxes', gt_bboxes[pos_gt_index]) return assign_result