Add checkpoints from the ablation study.

PiperOrigin-RevId: 328023346
This commit is contained in:
Florent Altché
2020-08-23 14:26:26 +01:00
committed by Diego de Las Casas
parent 22c3daff19
commit 8457046b2c
33 changed files with 397 additions and 363 deletions
+2 -2
View File
@@ -106,8 +106,8 @@ class Transporter(snt.AbstractModule):
num_keypoints = image_a_keypoints["heatmaps"].shape[-1]
transported_features = image_a_features
for k in range(num_keypoints):
mask_a = image_a_keypoints["heatmaps"][Ellipsis, k, None]
mask_b = image_b_keypoints["heatmaps"][Ellipsis, k, None]
mask_a = image_a_keypoints["heatmaps"][..., k, None]
mask_b = image_b_keypoints["heatmaps"][..., k, None]
# suppress features from image a, around both keypoint locations.
transported_features = (