def forward(self,
search,
template,
samp_idx=None,
labels=None,
settings=None):
""" The forward expects a NestedTensor, which consists of:
- search.tensors: batched images, of shape [batch_size x 3 x H_search x W_search]
- search.mask: a binary mask of shape [batch_size x H_search x W_search], containing 1 on padded pixels
- template.tensors: batched images, of shape [batch_size x 3 x H_template x W_template]
- template.mask: a binary mask of shape [batch_size x H_template x W_template], containing 1 on padded pixels
It returns a dict with the following elements:
- "pred_logits": the classification logits for all feature vectors.
Shape= [batch_size x num_vectors x (num_classes + 1)]
- "pred_boxes": The normalized boxes coordinates for all feature vectors, represented as
(center_x, center_y, height, width). These values are normalized in [0, 1],
relative to the size of each individual image.
"""
if not isinstance(search, NestedTensor):
search = nested_tensor_from_tensor(search)
if not isinstance(template, NestedTensor):
template = nested_tensor_from_tensor(template)
feature_search, pos_search = self.backbone(search)
feature_template, pos_template = self.backbone(template)
src_search, mask_search = feature_search[-1].decompose()
assert mask_search is not None
src_template, mask_template = feature_template[-1].decompose()
assert mask_template is not None
hs = self.featurefusion_network(self.input_proj(src_template),
mask_template,
self.input_proj(src_search),
mask_search, pos_template[-1],
pos_search[-1])
outputs_class = self.new_class_embed(hs)
outputs_coord = self.new_bbox_embed(hs).sigmoid()
out = {
'pred_logits': outputs_class[-1],
'pred_boxes': outputs_coord[-1]
}
return out
def forward(self, search, template, samp_idx, labels, settings, boxes):
""" The forward expects a NestedTensor, which consists of:
- search.tensors: batched images, of shape [batch_size x 3 x H_search x W_search]
- search.mask: a binary mask of shape [batch_size x H_search x W_search], containing 1 on padded pixels
- template.tensors: batched images, of shape [batch_size x 3 x H_template x W_template]
- template.mask: a binary mask of shape [batch_size x H_template x W_template], containing 1 on padded pixels
It returns a dict with the following elements:
- "pred_logits": the classification logits for all feature vectors.
Shape= [batch_size x num_vectors x (num_classes + 1)]
- "pred_boxes": The normalized boxes coordinates for all feature vectors, represented as
(center_x, center_y, height, width). These values are normalized in [0, 1],
relative to the size of each individual image.
"""
# Reshape search into a 4D tensor
search_shape = [int(x) for x in search.shape]
searcht = search.view([search_shape[0] * search_shape[1]] +
search_shape[2:])
if not isinstance(searcht, NestedTensor):
searcht = nested_tensor_from_tensor(searcht)
if not isinstance(template, NestedTensor):
templatet = nested_tensor_from_tensor(template)
with torch.no_grad():
feature_search, pos_search = self.backbone(searcht)
feature_template, pos_template = self.backbone(templatet)
src_search, mask_search = feature_search[-1].decompose()
assert mask_search is not None
src_template, mask_template = feature_template[-1].decompose()
assert mask_template is not None
circuit_input, _ = feature_search[-2].decompose()
post_src_search_shape = [int(x) for x in src_search.shape]
post_mask_search_shape = [int(x) for x in mask_search.shape]
src_search = src_search.view(search_shape[:2] +
post_src_search_shape[1:])
src_template = self.input_proj(src_template)
mask_search = mask_search.view(search_shape[:2] +
post_mask_search_shape[1:])
circuit_input = circuit_input.view(search_shape[0], search_shape[1],
-1, circuit_input.shape[2],
circuit_input.shape[3]).permute(
0, 2, 1, 3, 4)
proc_labels = F.interpolate(labels, circuit_input.shape[-2:])
proc_label_shape = [int(x) for x in proc_labels.shape]
proc_labels = proc_labels.view([search_shape[0]] +
proc_label_shape[1:]).to(
src_search.device) # .mean(2)
inh_1 = self.cnl(self.circuit_inh_1_init(1 - proc_labels))
exc_1 = self.cnl(self.circuit_exc_1_init(proc_labels))
nl_src_search = self.cnl(self.circuit_proj(circuit_input))
pos_search_shape = [int(x) for x in pos_search[-1].shape]
pos_search = pos_search[-1].view(search_shape[:2] +
pos_search_shape[1:])
mask_search = mask_search.view(search_shape[:2] +
post_mask_search_shape[1:])
excs = []
for t in range(nl_src_search.shape[2]):
# Step 1, saliency
pre_exc_1, inh_1 = self.circuit_1(nl_src_search[:, :, t],
excitation=exc_1,
inhibition=inh_1,
activ=self.cnl)
# Step 2, tracking
post_exc_1 = F.max_pool2d(self.cnl(
self.circuit_step2_trans(pre_exc_1)),
kernel_size=(2, 2),
stride=(2, 2),
padding=(0, 0))
if t == 0:
exc_2 = self.cnl(self.circuit_exc_2_init(post_exc_1))
inh_2 = self.cnl(self.circuit_inh_2_init(post_exc_1))
exc_2, inh_2 = self.circuit_2(
post_exc_1, excitation=exc_2, inhibition=inh_2,
activ=self.cnl) # , label=proc_labels)
# Split off a pair of TransT features and then use the circuit to gate the Qs in the decoder.
interp_exc_2 = F.interpolate(exc_2,
pre_exc_1.shape[2:],
mode="bilinear",
align_corners=False)
dec_rnn = self.circuit_rnn_decode_1(interp_exc_2)
prev_hs = self.circuit_rnn_decode_2(
self.nl(dec_rnn)) # * torch.sigmoid(self.rnn_gate(dec_rnn))
# src_search = src_search * proc_rnn.sigmoid()
# Pass activities through transformer
# hs, _ = self.featurefusion_network(src_template, mask_template, self.input_proj(src_search[:, t]), mask_search[:, t], pos_template[-1], pos_search[:, t], exc=prev_hs)
src_input = self.input_proj(src_search[:, t]) * prev_hs.sigmoid()
hs = self.featurefusion_network(src_template, mask_template,
src_input, mask_search[:, t],
pos_template[-1], pos_search[:, t])
# Step 3, TD-FB incorporating hs
res_hs = hs.squeeze().view(search_shape[0], self.height,
self.height,
self.hidden_dim).permute(0, 3, 1, 2)
# TD from Trans to Circuit
res_hs = F.max_pool2d(self.cnl(self.circuit_tf_codeswitch(res_hs)),
kernel_size=(2, 2),
stride=(2, 2),
padding=(0, 0))
##### This pooling is for technical reasons. Ideally dont have to do this but whatever...
if t == 0:
td_inh_2 = self.cnl(self.circuit_td_inh_2_init(res_hs))
exc_2, td_inh_2 = self.circuit_td_2(exc_2,
excitation=res_hs,
inhibition=td_inh_2,
activ=self.cnl)
# TD from circuit2 to Circuit 1
interp_exc_2 = self.cnl(
self.circuit_step3_trans(interp_exc_2)
) # F.interpolate(exc_2, pre_exc_1.shape[2:], mode="bilinear", align_corners=False)))
if t == 0:
td_inh_1 = self.cnl(self.circuit_td_inh_1_init(interp_exc_2))
exc_1, td_inh_1 = self.circuit_td_1(pre_exc_1,
excitation=interp_exc_2,
inhibition=td_inh_1,
activ=self.cnl)
# from matplotlib import pyplot as plt
# plt.subplot(141);plt.imshow(search[0, t].squeeze().permute(1, 2, 0).cpu());plt.subplot(142);
# plt.imshow((self.input_proj(src_search[:, t])[0].squeeze() ** 2).mean(0).detach().cpu());
# plt.subplot(143);plt.imshow((prev_hs.sigmoid().squeeze() ** 2)[0].mean(0).detach().cpu());
# plt.subplot(144);plt.imshow(((self.input_proj(src_search[:, t]) * prev_hs.sigmoid())[0].squeeze() ** 2).mean(0).detach().cpu());plt.show()
# if t > src_search.shape[1] - 5:
# plt.subplot(141);plt.title("Resnet");plt.imshow(search[0, t].squeeze().permute(1, 2, 0).cpu());plt.title("L1");plt.subplot(142);plt.imshow((pre_exc_1[0] ** 2).squeeze().mean(0).detach().cpu());plt.subplot(143);plt.title("L2");plt.imshow((exc_2[0] ** 2).squeeze().mean(0).detach().cpu());plt.subplot(144);plt.title("L1 after TD");plt.imshow((exc_1[0] ** 2).squeeze().mean(0).detach().cpu());
# plt.show()
excs.append(exc_1) # Also try exc_1?
# Concat exc to hs too
outputs_class = self.class_embed(hs)
outputs_coord = self.bbox_embed(hs).sigmoid()
excs = torch.stack(excs, 2)
proc_rnn = self.circuit_exc_bbox_3(
self.cnl(self.circuit_exc_bbox_1(excs)).view(
excs.shape[0], excs.shape[2], -1)).sigmoid()
if self.vj_pen:
out = {
'pred_logits': outputs_class[-1],
'pred_boxes': outputs_coord[-1],
'hgru_boxes': exc_bbox,
"vj_penalty": vj_penalty
}
else:
out = {
'pred_logits': outputs_class[-1],
'pred_boxes': outputs_coord[-1],
'hgru_boxes': proc_rnn
}
return out
def forward(self, search, template, samp_idx, labels, settings):
""" The forward expects a NestedTensor, which consists of:
- search.tensors: batched images, of shape [batch_size x 3 x H_search x W_search]
- search.mask: a binary mask of shape [batch_size x H_search x W_search], containing 1 on padded pixels
- template.tensors: batched images, of shape [batch_size x 3 x H_template x W_template]
- template.mask: a binary mask of shape [batch_size x H_template x W_template], containing 1 on padded pixels
It returns a dict with the following elements:
- "pred_logits": the classification logits for all feature vectors.
Shape= [batch_size x num_vectors x (num_classes + 1)]
- "pred_boxes": The normalized boxes coordinates for all feature vectors, represented as
(center_x, center_y, height, width). These values are normalized in [0, 1],
relative to the size of each individual image.
"""
# Reshape search into a 4D tensor
search_shape = [int(x) for x in search.shape]
search = search.view([search_shape[0] * search_shape[1]] + search_shape[2:])
if not isinstance(search, NestedTensor):
search = nested_tensor_from_tensor(search)
if not isinstance(template, NestedTensor):
template = nested_tensor_from_tensor(template)
with torch.no_grad():
feature_search, pos_search = self.backbone(search)
feature_template, pos_template = self.backbone(template)
src_search, mask_search = feature_search[-1].decompose()
assert mask_search is not None
src_template, mask_template = feature_template[-1].decompose()
assert mask_template is not None
# Use the circuit to track through the features
src_search = self.input_proj(src_search)
post_src_search_shape = [int(x) for x in src_search.shape]
post_mask_search_shape = [int(x) for x in mask_search.shape]
src_search = src_search.view(search_shape[:2] + post_src_search_shape[1:])
mask_search = mask_search.view(search_shape[:2] + post_mask_search_shape[1:])
proc_labels = self._generate_label_function(
labels,
settings.sigma,
settings.kernel,
settings.feature, # post_src_search_shape[-1], # settings.feature,
settings.output_sz, # post_src_search_shape[-1], # settings.output_sz,
settings.end_pad_if_even)
exc, inh = None, None
# pos_search_shape = [int(x) for x in pos_search[-1].shape]
proc_label_shape = [int(x) for x in proc_labels.shape]
# proc_labels = pos_search[-1].view(search_shape[:2] + pos_search_shape[1:]).mean(2)
proc_labels = proc_labels.view(search_shape[:2] + [1] + proc_label_shape[1:]).to(src_search.device) # .mean(2)
# from matplotlib import pyplot as plt
# im=2;ti=6;plt.subplot(121);plt.imshow(proc_labels[im, ti].cpu());plt.subplot(122);plt.imshow(src_search[im, ti].mean(0).detach().cpu());plt.show()
for t in range(samp_idx):
exc, inh = self.circuit(src_search[:, t], excitation=exc, inhibition=inh, label=proc_labels[:, t])
# Reshape exc for the transformers
exc = exc.flatten(2).permute(2, 0, 1)
# Reshape pos_search
pos_search_shape = [int(x) for x in pos_search[-1].shape]
pos_search = pos_search[-1].view(search_shape[:2] + pos_search_shape[1:])
# Split off a pair of TransT features and then use the circuit to gate the Qs in the decoder.
src_search = src_search[:, samp_idx]
mask_search = mask_search[:, samp_idx]
pos_search = pos_search[:, samp_idx]
hs = self.featurefusion_network(self.input_proj(src_template), mask_template, src_search, mask_search, pos_template[-1], pos_search, exc=exc)
# Concat exc to hs too
outputs_class = self.class_embed(hs)
outputs_coord = self.bbox_embed(hs).sigmoid()
out = {'pred_logits': outputs_class[-1], 'pred_boxes': outputs_coord[-1]}
return out
def forward(self, search, template, samp_idx, labels, settings, boxes):
""" The forward expects a NestedTensor, which consists of:
- search.tensors: batched images, of shape [batch_size x 3 x H_search x W_search]
- search.mask: a binary mask of shape [batch_size x H_search x W_search], containing 1 on padded pixels
- template.tensors: batched images, of shape [batch_size x 3 x H_template x W_template]
- template.mask: a binary mask of shape [batch_size x H_template x W_template], containing 1 on padded pixels
It returns a dict with the following elements:
- "pred_logits": the classification logits for all feature vectors.
Shape= [batch_size x num_vectors x (num_classes + 1)]
- "pred_boxes": The normalized boxes coordinates for all feature vectors, represented as
(center_x, center_y, height, width). These values are normalized in [0, 1],
relative to the size of each individual image.
"""
# Reshape search into a 4D tensor
search_shape = [int(x) for x in search.shape]
searcht = search.view([search_shape[0] * search_shape[1]] + search_shape[2:])
if not isinstance(searcht, NestedTensor):
searcht = nested_tensor_from_tensor(searcht)
if not isinstance(template, NestedTensor):
templatet = nested_tensor_from_tensor(template)
with torch.no_grad():
feature_search, pos_search = self.backbone(searcht)
feature_template, pos_template = self.backbone(templatet)
src_search, mask_search = feature_search[-1].decompose()
assert mask_search is not None
src_template, mask_template = feature_template[-1].decompose()
assert mask_template is not None
circuit_input, _ = feature_search[-1].decompose()
post_src_search_shape = [int(x) for x in src_search.shape]
post_mask_search_shape = [int(x) for x in mask_search.shape]
src_search = src_search.view(search_shape[:2] + post_src_search_shape[1:])
src_template = self.input_proj(src_template)
mask_search = mask_search.view(search_shape[:2] + post_mask_search_shape[1:])
circuit_input = circuit_input.view(search_shape[0], search_shape[1], -1, circuit_input.shape[2], circuit_input.shape[3]).permute(0, 2, 1, 3, 4)
proc_labels = F.interpolate(labels, circuit_input.shape[-2:])
proc_label_shape = [int(x) for x in proc_labels.shape]
proc_labels = proc_labels.view([search_shape[0]] + proc_label_shape[1:]).to(src_search.device) # .mean(2)
inh_1 = self.cnl(self.circuit_inh_1_init(1 - proc_labels))
exc_1 = self.cnl(self.circuit_exc_1_init(proc_labels))
nl_src_search = self.cnl(self.circuit_proj(circuit_input))
pos_search_shape = [int(x) for x in pos_search[-1].shape]
pos_search = pos_search[-1].view(search_shape[:2] + pos_search_shape[1:])
mask_search = mask_search.view(search_shape[:2] + post_mask_search_shape[1:])
excs, rnn_gates = [], []
for t in range(nl_src_search.shape[2]):
# Step 1, saliency
pre_exc_1, inh_1 = self.circuit_1(nl_src_search[:, :, t], excitation=exc_1, inhibition=inh_1, activ=self.cnl)
# Split off a pair of TransT features and then use the circuit to gate the Qs in the decoder.
dec_rnn = self.circuit_rnn_decode_1(pre_exc_1)
prev_hs = self.circuit_rnn_decode_2(self.cnl(dec_rnn)) # * torch.sigmoid(self.rnn_gate(dec_rnn))
# Pass activities through transformer
# prev_hs = 1 + (prev_hs.tanh()) # in [0, 2]
if t > 0:
# cost_vol = torch.exp(-(res_hs - pre_exc_1) ** 2) # Correspondence between FB and the predicted modulation
# This cost vol is all-to-all spatial maps. Then Reshaped so that all features are aggregated.
cost_vol = torch.einsum('bik,bjk->bijk', res_hs.view(pre_exc_1.shape[0], pre_exc_1.shape[1], -1), pre_exc_1.view(pre_exc_1.shape[0], pre_exc_1.shape[1], -1)).view(pre_exc_1.shape[0], pre_exc_1.shape[1] ** 2, pre_exc_1.shape[2], pre_exc_1.shape[3])
# cost_vol = torch.cat([res_hs, prev_hs], 1)
rnn_gate = self.circuit_gate(cost_vol).sigmoid()
# rnn_gates.append((rnn_gate ** 2).mean((1, 2, 3)))
src_input = self.input_proj(src_search[:, t]) + prev_hs * rnn_gate
else:
rnn_gate = 0. # Dont allow the RNN on the first frame
src_input = self.input_proj(src_search[:, t])
hs = self.featurefusion_network(src_template, mask_template, src_input, mask_search[:, t], pos_template[-1], pos_search[:, t])
# Step 3, TD-FB incorporating hs
pre_res_hs = hs.squeeze().view(search_shape[0], self.height, self.height, self.hidden_dim).permute(0, 3, 1, 2)
# TD from Trans to Circuit
res_hs = self.cnl(self.circuit_tf_codeswitch(pre_res_hs))
##### This pooling is for technical reasons. Ideally dont have to do this but whatever...
if t == 0:
td_inh_1 = self.cnl(self.circuit_td_inh_1_init(res_hs))
exc_1, td_inh_1 = self.circuit_td_1(pre_exc_1, excitation=res_hs, inhibition=td_inh_1, activ=self.cnl)
# from matplotlib import pyplot as plt
# plt.subplot(141);plt.imshow(search[0, t].squeeze().permute(1, 2, 0).cpu());plt.subplot(142);
# plt.imshow((self.input_proj(src_search[:, t])[0].squeeze() ** 2).mean(0).detach().cpu());
# plt.subplot(143);plt.imshow((prev_hs.sigmoid().squeeze() ** 2)[0].mean(0).detach().cpu());
# plt.subplot(144);plt.imshow(((self.input_proj(src_search[:, t]) * prev_hs.sigmoid())[0].squeeze() ** 2).mean(0).detach().cpu());plt.show()
# if t > src_search.shape[1] - 5:
# plt.subplot(141);plt.title("Resnet");plt.imshow(search[0, t].squeeze().permute(1, 2, 0).cpu());plt.title("L1");plt.subplot(142);plt.imshow((pre_exc_1[0] ** 2).squeeze().mean(0).detach().cpu());plt.subplot(143);plt.title("L2");plt.imshow((exc_2[0] ** 2).squeeze().mean(0).detach().cpu());plt.subplot(144);plt.title("L1 after TD");plt.imshow((exc_1[0] ** 2).squeeze().mean(0).detach().cpu());
# plt.show()
excs.append(exc_1)
# Concat exc to hs too
outputs_class = self.class_embed(hs)
outputs_coord = self.bbox_embed(hs).sigmoid()
excs = torch.stack(excs, 2)
proc_rnn = self.circuit_exc_bbox_3(self.cnl(self.circuit_exc_bbox_1(excs)).view(excs.shape[0], excs.shape[2], -1)).sigmoid()
rnn_gate = self.circuit_gate_readout(self.cnl(cost_vol)).mean((1, 2, 3)) # [:, None]
out = {'pred_logits': outputs_class[-1], 'pred_boxes': outputs_coord[-1], 'hgru_boxes': proc_rnn, 'rnn_gate': rnn_gate}
return out
def forward(self, search, template, samp_idx, labels, settings, boxes):
""" The forward expects a NestedTensor, which consists of:
- search.tensors: batched images, of shape [batch_size x 3 x H_search x W_search]
- search.mask: a binary mask of shape [batch_size x H_search x W_search], containing 1 on padded pixels
- template.tensors: batched images, of shape [batch_size x 3 x H_template x W_template]
- template.mask: a binary mask of shape [batch_size x H_template x W_template], containing 1 on padded pixels
It returns a dict with the following elements:
- "pred_logits": the classification logits for all feature vectors.
Shape= [batch_size x num_vectors x (num_classes + 1)]
- "pred_boxes": The normalized boxes coordinates for all feature vectors, represented as
(center_x, center_y, height, width). These values are normalized in [0, 1],
relative to the size of each individual image.
"""
# Reshape search into a 4D tensor
search_shape = [int(x) for x in search.shape]
searcht = search.view([search_shape[0] * search_shape[1]] +
search_shape[2:])
if not isinstance(searcht, NestedTensor):
searcht = nested_tensor_from_tensor(searcht)
if not isinstance(template, NestedTensor):
templatet = nested_tensor_from_tensor(template)
with torch.no_grad():
feature_search, pos_search = self.backbone(searcht)
feature_template, pos_template = self.backbone(templatet)
src_search, mask_search = feature_search[-1].decompose()
assert mask_search is not None
src_template, mask_template = feature_template[-1].decompose()
assert mask_template is not None
post_src_search_shape = [int(x) for x in src_search.shape]
post_mask_search_shape = [int(x) for x in mask_search.shape]
src_search = src_search.view(search_shape[:2] +
post_src_search_shape[1:])
mask_search = mask_search.view(search_shape[:2] +
post_mask_search_shape[1:])
# Reshape pos_search
pos_search_shape = [int(x) for x in pos_search[-1].shape]
pos_search = pos_search[-1].view(search_shape[:2] +
pos_search_shape[1:])
# Split off a pair of TransT features and then use the circuit to gate the Qs in the decoder.
src_search = self.input_proj(src_search[:, -1]) # samp_idx
src_template = self.input_proj(src_template)
mask_search = mask_search[:, -1] # samp_idx]
pos_search = pos_search[:, -1] # samp_idx]
hs = self.featurefusion_network(
src_template, mask_template, src_search, mask_search,
pos_template[-1],
pos_search) # , exc=proc_rnn[:, :, -1]) # .sigmoid())
# Concat exc to hs too
# hs = torch.cat([hs, self.rnn_embed(exc)], -1)
outputs_class = self.class_embed(hs)
outputs_coord = self.bbox_embed(hs).sigmoid()
out = {
'pred_logits': outputs_class[-1],
'pred_boxes': outputs_coord[-1],
'hgru_boxes': src_search
}
return out
def forward(self, search, template, samp_idx, labels, settings, boxes):
""" The forward expects a NestedTensor, which consists of:
- search.tensors: batched images, of shape [batch_size x 3 x H_search x W_search]
- search.mask: a binary mask of shape [batch_size x H_search x W_search], containing 1 on padded pixels
- template.tensors: batched images, of shape [batch_size x 3 x H_template x W_template]
- template.mask: a binary mask of shape [batch_size x H_template x W_template], containing 1 on padded pixels
It returns a dict with the following elements:
- "pred_logits": the classification logits for all feature vectors.
Shape= [batch_size x num_vectors x (num_classes + 1)]
- "pred_boxes": The normalized boxes coordinates for all feature vectors, represented as
(center_x, center_y, height, width). These values are normalized in [0, 1],
relative to the size of each individual image.
"""
# Reshape search into a 4D tensor
search_shape = [int(x) for x in search.shape]
searcht = search.view([search_shape[0] * search_shape[1]] +
search_shape[2:])
if not isinstance(searcht, NestedTensor):
searcht = nested_tensor_from_tensor(searcht)
if not isinstance(template, NestedTensor):
templatet = nested_tensor_from_tensor(template)
with torch.no_grad():
feature_search, pos_search = self.backbone(searcht)
feature_template, pos_template = self.backbone(templatet)
src_search, mask_search = feature_search[-1].decompose()
assert mask_search is not None
src_template, mask_template = feature_template[-1].decompose()
assert mask_template is not None
post_src_search_shape = [int(x) for x in src_search.shape]
post_mask_search_shape = [int(x) for x in mask_search.shape]
src_search = src_search.view(search_shape[:2] +
post_src_search_shape[1:])
mask_search = mask_search.view(search_shape[:2] +
post_mask_search_shape[1:])
proc_labels = F.interpolate(labels, src_search.shape[-2:])
# Duble check that the below is correct
# rnn_src_search = self.nl(self.rnn_proj(src_search.permute(0, 2, 1, 3, 4)))
# psrc = src_search.permute(0, 2, 1, 3, 4)
# rnn_channels = self.rnn_proj_channel_2(self.nl(self.rnn_proj_channel_1(psrc.mean(dim=(3, 4), keepdim=True))))
# rnn_space = self.rnn_proj_spatial_2(self.nl(self.rnn_proj_spatial_1(psrc)))
# rnn_src_search = F.softplus(rnn_space + rnn_channels) # Used to be multiplicative
proc_label_shape = [int(x) for x in proc_labels.shape]
proc_labels = proc_labels.view([search_shape[0]] +
proc_label_shape[1:]).to(
src_search.device) # .mean(2)
exc_1, inh_1 = None, None
exc_2, inh_2 = None, None
td_inh = None
excs = []
# for t in range(rnn_src_search.shape[2]):
# inh_1 = self.nl(self.inh_1_proj(src_search[:, 0]))
inh_1 = self.nl(self.inh_1_init(src_search[:, 0]))
exc_1 = self.nl(self.exc_1_init(src_search[:, 0]))
for t in range(src_search.shape[1]):
# Step 1, saliency
pre_exc_1, inh_1 = self.circuit_1(self.nl(src_search[:, t]),
excitation=exc_1,
inhibition=inh_1)
# Step 2, tracking
exc_2, inh_2 = self.circuit_2(self.step2_trans(pre_exc_1),
excitation=exc_2,
inhibition=inh_2,
label=proc_labels)
# Step 3, TD-FB
if t == 0:
# td_inh = self.nl(self.td_inh_proj(exc_2))
td_inh = self.nl(torch.zeros_like(exc_2))
td_exc = self.nl(self.step3_trans(exc_2))
# exc_1, td_inh = self.circuit_td(pre_exc_1, excitation=self.nl(self.step3_trans(exc_2)), inhibition=td_inh)
exc_1, td_inh = self.circuit_td(pre_exc_1,
excitation=td_exc,
inhibition=td_inh)
# plt.subplot(141);plt.imshow(search[0, t].squeeze().permute(1, 2, 0).cpu());plt.subplot(142);plt.imshow((rnn_src_search[0, :, t].squeeze() ** 2).mean(0).detach().cpu());plt.subplot(143);plt.imshow((exc[0].squeeze() ** 2).mean(0).detach().cpu());plt.subplot(144);plt.imshow((inh[0].squeeze() ** 2).mean(0).detach().cpu());plt.show()
plt.subplot(141)
plt.imshow(search[0, t].squeeze().permute(1, 2, 0).cpu())
plt.subplot(142)
plt.imshow((pre_exc_1[0]**2).squeeze().mean(0).detach().cpu())
plt.subplot(143)
plt.imshow((exc_2[0]**2).squeeze().mean(0).detach().cpu())
plt.subplot(144)
plt.imshow((exc_1[0]**2).squeeze().mean(0).detach().cpu())
# plt.show()
excs.append(exc_1) # Also try exc_1?
# ros.append(self.rnn_embed(exc_2.flatten(2).permute(2, 0, 1)).sigmoid())
# if t == samp_idx - 2 and self.vj_pen:
# penultimate = exc.clone()
# ros = self.rnn_embed(exc_2.flatten(2).permute(2, 0, 1).unsqueeze(0).transpose(1, 2)).sigmoid()
if self.vj_pen:
norm_1_vect = torch.ones_like(exc)
norm_1_vect.requires_grad = False
vj_prod = torch.autograd.grad(exc,
penultimate,
grad_outputs=[norm_1_vect],
retain_graph=True,
create_graph=True,
allow_unused=True)[0]
vj_penalty = (vj_prod -
0.95).clamp(0)**2 # Squared to emphasize outliers
vj_penalty = vj_penalty.sum() # Save memory with the mean
# from matplotlib import pyplot as plt
# im=0;ti=0;plt.subplot(141);plt.imshow(proc_labels[im].squeeze().cpu());plt.subplot(142);plt.imshow((rnn_src_search[im, :, ti] ** 2).mean(0).detach().cpu());plt.subplot(143);plt.imshow((exc[im] ** 2).squeeze().mean(0).detach().cpu());plt.subplot(144);plt.imshow(search[im, ti].squeeze().permute(1, 2, 0).cpu());plt.show()
# Reshape pos_search
pos_search_shape = [int(x) for x in pos_search[-1].shape]
pos_search = pos_search[-1].view(search_shape[:2] +
pos_search_shape[1:])
# Split off a pair of TransT features and then use the circuit to gate the Qs in the decoder.
src_search = self.input_proj(src_search[:, -1]) # samp_idx
src_template = self.input_proj(src_template)
mask_search = mask_search[:, -1] # samp_idx]
pos_search = pos_search[:, -1] # samp_idx]
dec_rnn = self.rnn_decode_1(self.bn(torch.stack(excs, 2)))
proc_rnn = self.rnn_decode_2(
self.nl(dec_rnn)) # * torch.sigmoid(self.rnn_gate(dec_rnn))
hs, _ = self.featurefusion_network(src_template,
mask_template,
src_search,
mask_search,
pos_template[-1],
pos_search,
exc=proc_rnn[:, :, -1])
# Concat exc to hs too
# hs = torch.cat([hs, self.rnn_embed(exc)], -1)
outputs_class = self.class_embed(hs)
outputs_coord = self.bbox_embed(hs).sigmoid()
# exc_bbox = self.exc_bbox(exc).sigmoid()
# coords = torch.stack(torch.meshgrid(torch.arange(excs.shape[-2]), torch.arange(excs.shape[-1])), 0)
# coords = coords / coords.max()
# coords = coords.float().to(excs.device)[None, :, None].repeat(excs.shape[0], 1, excs.shape[2], 1, 1)
# excs = torch.cat([coords, excs], 1)e
# exc = exc.view(exc.shape[1], self.height, self.height, -1).permute(0, 3, 1, 2)
# excs = self.exc_bbox_2(F.relu(self.exc_bbox_1(excs)))
# exc_bbox = self.exc_bbox_3(excs.flatten(3).squeeze(1)).sigmoid()
proc_rnn = self.exc_bbox_3(
self.exc_bbox_1(proc_rnn).view(proc_rnn.shape[0],
proc_rnn.shape[2], -1)).sigmoid()
if self.vj_pen:
out = {
'pred_logits': outputs_class[-1],
'pred_boxes': outputs_coord[-1],
'hgru_boxes': exc_bbox,
"vj_penalty": vj_penalty
}
else:
out = {
'pred_logits': outputs_class[-1],
'pred_boxes': outputs_coord[-1],
'hgru_boxes': proc_rnn
}
return out