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
