def training(self, head, x, proposals, bbx, cat, iscrowd, ids, msk,
img_size):
x = x[self.min_level:self.min_level + self.levels]
try:
if proposals.all_none:
raise Empty
# Match proposals to ground truth
with torch.no_grad():
proposals, match = self.proposal_matcher(
proposals, bbx, cat, iscrowd)
cls_lbl, bbx_lbl, msk_lbl = self._match_to_lbl(
proposals, bbx, cat, ids, msk, match)
if proposals.all_none:
raise Empty
# Run head
set_active_group(head, active_group(True))
proposals, proposals_idx = proposals.contiguous
cls_logits, bbx_logits, msk_logits = self._head(
head, x, proposals, proposals_idx, img_size, True, True)
# Calculate losses
cls_loss, bbx_loss = self.bbx_loss(cls_logits, bbx_logits, cls_lbl,
bbx_lbl)
msk_loss = self.msk_loss(msk_logits, cls_lbl, msk_lbl)
except Empty:
active_group(False)
cls_loss = bbx_loss = msk_loss = sum(x_i.sum() for x_i in x) * 0
return cls_loss, bbx_loss, msk_loss
def training(self, head, x, labels, img_size):
x = self._get_level(x)
try:
# Run head
set_active_group(head, active_group(True))
labels, labels_idx = labels.contiguous
ret_pred = self._head(head, x)
# Calculate losses
ret_loss = self.loss(ret_pred, labels, labels_idx)
except Empty:
active_group(False)
ret_loss = sum(x_i.sum() for x_i in x) * 0
return ret_loss, ret_pred
def forward(self, head_coarse, head_fine, x1, x2, kpts, scores, img_size_1,
img_size_2, valid_size_1, valid_size_2, intrinsics1,
intrinsics2, extrinsics1, extrinsics2):
# select fpn levels
x1_f, x1_c = self._get_level(x1)
x2_f, x2_c = self._get_level(x2)
data = {
'B': x1[0].size(0),
'img_size_1': img_size_1,
'img_size_2': img_size_2,
'coarse_size_1': x1_c.shape[2:],
'coarse_size_2': x2_c.shape[2:],
'fine_size_1': x1_f.shape[2:],
'fine_size_2': x2_f.shape[2:],
'dim_c': head_coarse.embedding_size,
'dim_f': head_fine.embedding_size,
'win_size': self.win_size
}
try:
# 1. sparse annotation generation
if kpts.all_none:
raise Empty
with torch.no_grad():
# Compute F and P matrices
P, E, F = self.get_transformation(intrinsics1, intrinsics2,
extrinsics1, extrinsics2)
# Keypoint Generator
kpts = self.generator(kpts, scores, valid_size_2, intrinsics1,
intrinsics2, F, P)
if kpts.all_none:
raise Empty
# Get contiguious view and valid transformations
kpts, kpts_idx = kpts.contiguous
# 2. coarse-level module
# add featmap with positional encoding, then flatten it to sequence [N, HW, C]
desc_c1 = rearrange(self.pos_encoding(x1_c),
'n c h w -> n (h w) c')
desc_c2 = rearrange(self.pos_encoding(x2_c),
'n c h w -> n (h w) c')
mask_c0 = mask_c1 = None # mask is useful in training
# if 'mask0' in data:
# mask_c0, mask_c1 = data['mask0'].flatten(-2), data['mask1'].flatten(-2)
# Run coarse head
set_active_group(head_coarse, active_group(True))
desc_c1, desc_c2 = head_coarse(desc_c1, desc_c2, mask_c0, mask_c1)
# 3. match coarse-level
matches = self.coarse_matching(desc_c1,
desc_c2,
data,
mask_c0=mask_c1,
mask_c1=mask_c1)
# 4. coarse to fine refinement
# feat_f0_unfold, feat_f1_unfold = self.fine_preprocess(feat_f0, feat_f1, feat_c0, feat_c1, data)
desc_f1, desc_f2 = self.get_fine_windows_prediction(
desc_c1, desc_c2, x1_f, x2_f, matches, data)
# 5. fine-level module
if desc_f1.size(0) != 0: # at least one coarse level predicted
# Run fine head
set_active_group(head_fine, active_group(True))
desc_f1, desc_f2 = head_fine(desc_f1, desc_f2)
# 6. match fine-level
matches = self.fine_matching(desc_f1, desc_f2, matches, data)
# Compute Loss
F, _ = F.contiguous
P, _ = P.contiguous
epipolar_loss = self.loss(kpts, F[kpts_idx], P[kpts_idx], matches,
data)
except Empty:
active_group(False)
epipolar_loss = sum(x_i.sum() for x_i in x1) * 0
return epipolar_loss
def training(self, head, x, rpn_proposals, bbx, cat, iscrowd, ids, msk,
img_size, valid_size):
x = x[self.min_level:self.min_level + self.levels]
cls_loss_sum, bbx_loss_sum, msk_loss_sum = 0.0, 0.0, 0.0
y = None
try:
#### Zero pooling ###
if rpn_proposals.all_none:
raise Empty
#Match proposals to ground truth
with torch.no_grad():
# get stage proposals Top 512 samples
proposals, match, rpn_proposals, indices = self.proposal_matcher(
rpn_proposals, bbx, cat, iscrowd, 0)
#np.savetxt('rpn_proposals_original', rpn_proposals[0].cpu(), fmt='%1.2f')
cls_lbl, bbx_lbl, msk_lbl = self._match_to_lbl(
proposals, bbx, cat, ids, msk, match)
if proposals.all_none:
raise Empty
_proposals, _proposals_idx = proposals.contiguous
for stage in range(len(self.stage_loss_weights)):
#print('__ {} __'.format(stage))
# Current head and its weighted loss coefficient lw
current_head = head[stage]
lw = self.stage_loss_weights[stage]
### Run Current head for BBox training & loss calculation ###
set_active_group(current_head, active_group(True))
cls_logits, bbx_logits, _, _ = self._head(
current_head, x, y, _proposals, _proposals_idx, img_size,
True, False)
cls_loss, bbx_loss = self.bbx_loss(cls_logits, bbx_logits,
cls_lbl, bbx_lbl)
cls_loss_sum += lw * cls_loss
bbx_loss_sum += lw * bbx_loss
#print('cls_loss_{} ='.format(stage), cls_loss.item())
#print('bbx_loss_{} ='.format(stage), bbx_loss.item())
### Get new proposals for next stage ###
_proposals, _indices = self._inference(x, _proposals,
_proposals_idx,
bbx_logits, cls_logits,
valid_size, indices,
stage)
if _proposals.all_none:
raise Empty
# Flip to replace repeated indices with best scoring bbox
rpn_proposals_augument = []
for _proposals_i, _indices_i, rpn_proposals_i in zip(
_proposals, _indices, rpn_proposals):
rpn_proposals_i[_indices_i.flip(0)] = _proposals_i.flip(0)
rpn_proposals_augument.append(rpn_proposals_i)
### New proposals for next stage ###
rpn_proposals = PackedSequence(rpn_proposals_augument)
if rpn_proposals.all_none:
raise Empty
#with torch.no_grad():
#np.savetxt('rpn_proposals_augument_{}.txt'.format(stage), rpn_proposals_augument[0].cpu().numpy(),
# fmt='%1.2f')
#### Stage pooling ###
with torch.no_grad():
# get stage proposals Top 512 samples
proposals, match, rpn_proposals, indices = self.proposal_matcher(
rpn_proposals, bbx, cat, iscrowd, stage)
#np.savetxt('rpn_proposals_{}'.format(stage), rpn_proposals[0].cpu(), fmt='%1.2f')
cls_lbl, bbx_lbl, msk_lbl = self._match_to_lbl(
proposals, bbx, cat, ids, msk, match)
# Run Previous head for Mask training & loss calculation
_proposals, _proposals_idx = _proposals.contiguous
set_active_group(current_head, active_group(True))
_, _, msk_logits, y = self._head(current_head, x, y,
_proposals, _proposals_idx,
img_size, False, True)
msk_loss = self.msk_loss(msk_logits, cls_lbl, msk_lbl)
msk_loss_sum += lw * msk_loss
#print('msk_loss_{} ='.format(stage), msk_loss.item())
except Empty:
active_group(False)
cls_loss_sum = bbx_loss_sum = msk_loss_sum = sum(x_i.sum()
for x_i in x) * 0
return cls_loss_sum, bbx_loss_sum, msk_loss_sum