def localize_target(self, scores, sample_pos, sample_scales):
"""Run the target localization."""
scores = scores.squeeze(1)
preprocess_method = self.params.get('score_preprocess', 'none')
if preprocess_method == 'none':
pass
elif preprocess_method == 'exp':
scores = scores.exp()
elif preprocess_method == 'softmax':
reg_val = getattr(self.net.classifier.filter_optimizer,
'softmax_reg', None)
scores_view = scores.view(scores.shape[0], -1)
scores_softmax = activation.softmax_reg(scores_view,
dim=-1,
reg=reg_val)
scores = scores_softmax.view(scores.shape)
else:
raise Exception('Unknown score_preprocess in params.')
score_filter_ksz = self.params.get('score_filter_ksz', 1)
if score_filter_ksz > 1:
assert score_filter_ksz % 2 == 1
kernel = scores.new_ones(1, 1, score_filter_ksz, score_filter_ksz)
scores = F.conv2d(scores.view(-1, 1, *scores.shape[-2:]),
kernel,
padding=score_filter_ksz // 2).view(scores.shape)
if self.params.get('advanced_localization', False):
return self.localize_advanced(scores, sample_pos, sample_scales)
# Get maximum
score_sz = torch.Tensor(list(scores.shape[-2:]))
score_center = (score_sz - 1) / 2
max_score, max_disp = dcf.max2d(scores)
_, scale_ind = torch.max(max_score, dim=0)
max_disp = max_disp[scale_ind, ...].float().cpu().view(-1)
target_disp = max_disp - score_center
# Compute translation vector and scale change factor
output_sz = score_sz - (self.kernel_size + 1) % 2
translation_vec = target_disp * (self.img_support_sz /
output_sz) * sample_scales[scale_ind]
return translation_vec, scale_ind, scores, None
def forward(self, weights, feat, bb, sample_weight=None, num_iter=None, compute_losses=True):
"""Runs the optimizer module.
Note that [] denotes an optional dimension.
args:
weights: Initial weights. Dims (sequences, feat_dim, wH, wW).
feat: Input feature maps. Dims (images_in_sequence, [sequences], feat_dim, H, W).
bb: Target bounding boxes (x, y, w, h) in the image coords. Dims (images_in_sequence, [sequences], 4).
sample_weight: Optional weight for each sample. Dims: (images_in_sequence, [sequences]).
num_iter: Number of iterations to run.
compute_losses: Whether to compute the (train) loss in each iteration.
returns:
weights: The final oprimized weights.
weight_iterates: The weights computed in each iteration (including initial input and final output).
losses: Train losses."""
# Sizes
num_iter = self.num_iter if num_iter is None else num_iter
num_images = feat.shape[0]
num_sequences = feat.shape[1] if feat.dim() == 5 else 1
filter_sz = (weights.shape[-2], weights.shape[-1])
output_sz = (feat.shape[-2] + (weights.shape[-2] + 1) % 2, feat.shape[-1] + (weights.shape[-1] + 1) % 2)
# Get learnable scalars
step_length_factor = torch.exp(self.log_step_length)
reg_weight = (self.filter_reg*self.filter_reg).clamp(min=self.min_filter_reg**2)
# Compute label density
offset = (torch.Tensor(filter_sz).to(bb.device) % 2) / 2.0
center = ((bb[..., :2] + bb[..., 2:] / 2) / self.feat_stride).flip((-1,)) - offset
label_density = self.get_label_density(center, output_sz)
# Get total sample weights
if sample_weight is None:
sample_weight = torch.Tensor([1.0 / num_images]).to(feat.device)
elif isinstance(sample_weight, torch.Tensor):
sample_weight = sample_weight.reshape(num_images, num_sequences, 1, 1)
exp_reg = 0 if self.softmax_reg is None else math.exp(self.softmax_reg)
def _compute_loss(scores, weights):
return torch.sum(sample_weight.reshape(sample_weight.shape[0], -1) *
(torch.log(scores.exp().sum(dim=(-2, -1)) + exp_reg) - (label_density * scores).sum(dim=(-2, -1)))) / num_sequences +\
reg_weight * (weights ** 2).sum() / num_sequences
weight_iterates = [weights]
losses = []
for i in range(num_iter):
if i > 0 and i % self.detach_length == 0:
weights = weights.detach()
# Compute "residuals"
scores = filter_layer.apply_filter(feat, weights)
scores_softmax = activation.softmax_reg(scores.reshape(num_images, num_sequences, -1), dim=2, reg=self.softmax_reg).reshape(scores.shape)
res = sample_weight*(scores_softmax - label_density)
if compute_losses:
losses.append(_compute_loss(scores, weights))
# Compute gradient
weights_grad = filter_layer.apply_feat_transpose(feat, res, filter_sz, training=self.training) + \
reg_weight * weights
# Map the gradient with the Hessian
scores_grad = filter_layer.apply_filter(feat, weights_grad)
sm_scores_grad = scores_softmax * scores_grad
hes_scores_grad = sm_scores_grad - scores_softmax * torch.sum(sm_scores_grad, dim=(-2,-1), keepdim=True)
grad_hes_grad = (scores_grad * hes_scores_grad).reshape(num_images, num_sequences, -1).sum(dim=2).clamp(min=0)
grad_hes_grad = (sample_weight.reshape(sample_weight.shape[0], -1) * grad_hes_grad).sum(dim=0)
# Compute optimal step length
alpha_num = (weights_grad * weights_grad).sum(dim=(1,2,3))
alpha_den = (grad_hes_grad + (reg_weight + self.alpha_eps) * alpha_num).clamp(1e-8)
alpha = alpha_num / alpha_den
# Update filter
weights = weights - (step_length_factor * alpha.reshape(-1, 1, 1, 1)) * weights_grad
# Add the weight iterate
weight_iterates.append(weights)
if compute_losses:
scores = filter_layer.apply_filter(feat, weights)
losses.append(_compute_loss(scores, weights))
return weights, weight_iterates, losses
def forward(self, meta_parameter: TensorList, num_iter=None, **kwargs):
if not isinstance(meta_parameter, TensorList):
meta_parameter = TensorList([meta_parameter])
_residual_batch_dim = 1
# Make sure grad is enabled
torch_grad_enabled = torch.is_grad_enabled()
torch.set_grad_enabled(True)
num_iter = self.num_iter if num_iter is None else num_iter
step_length_factor = torch.exp(self.log_step_length)
label_density, sample_weight, reg_weight = self.score_predictor.init_data(
meta_parameter, **kwargs)
exp_reg = 0 if self.softmax_reg is None else math.exp(self.softmax_reg)
def _compute_loss(scores, weights):
num_sequences = scores.shape[_residual_batch_dim]
return torch.sum(sample_weight.reshape(sample_weight.shape[0], -1) *
(torch.log(scores.exp().sum(dim=(-2, -1)) + exp_reg) - (label_density * scores).sum(dim=(-2, -1)))) / num_sequences + \
reg_weight * sum((weights * weights).sum()) / num_sequences
meta_parameter_iterates = [meta_parameter]
losses = []
for i in range(num_iter):
if i > 0 and i % self.detach_length == 0:
meta_parameter = meta_parameter.detach()
meta_parameter.requires_grad_(True)
# Compute residual vector
scores = self.score_predictor(meta_parameter, **kwargs)
if self.compute_losses:
losses.append(_compute_loss(scores, meta_parameter))
scores_softmax = activation.softmax_reg(
scores.reshape(*scores.shape[:2], -1),
dim=2,
reg=self.softmax_reg).reshape(scores.shape)
dLds = sample_weight * (scores_softmax - label_density)
# Compute gradient of loss
weights_grad = TensorList(torch.autograd.grad(scores, meta_parameter, dLds, create_graph=True)) + \
meta_parameter * reg_weight
# Multiply gradient with Jacobian
scores_grad = torch.autograd.grad(weights_grad,
dLds,
weights_grad,
create_graph=True)[0]
sm_scores_grad = scores_softmax * scores_grad
hes_scores_grad = sm_scores_grad - scores_softmax * torch.sum(sm_scores_grad, dim=(-2, -1), keepdim=True) + \
self.hessian_reg * scores_grad
grad_hes_grad = (scores_grad * hes_scores_grad).reshape(
*scores.shape[:2], -1).sum(dim=2).clamp(min=0)
grad_hes_grad = (
sample_weight.reshape(sample_weight.shape[0], -1) *
grad_hes_grad).sum(dim=0)
# Compute optimal step length
gg = (weights_grad * weights_grad).reshape(scores.shape[1],
-1).sum(dim=1)
alpha_num = sum(gg)
alpha_den = (grad_hes_grad + sum(gg * reg_weight) +
self.steplength_reg * alpha_num).clamp(1e-8)
alpha = step_length_factor * (alpha_num / alpha_den)
# Compute optimization step
step = weights_grad.apply(lambda e: alpha.reshape([
-1 if d == self._parameter_batch_dim else 1
for d in range(e.dim())
]) * e)
# Add step to parameter
meta_parameter = meta_parameter - step
meta_parameter_iterates.append(meta_parameter)
if self.compute_losses:
losses.append(
_compute_loss(self.score_predictor(meta_parameter, **kwargs),
meta_parameter))
# Reset the grad enabled flag
torch.set_grad_enabled(torch_grad_enabled)
if not torch_grad_enabled:
meta_parameter.detach_()
for w in meta_parameter_iterates:
w.detach_()
for l in losses:
l.detach_()
return meta_parameter, meta_parameter_iterates, losses