summaryrefslogtreecommitdiff
path: root/models
diff options
context:
space:
mode:
authorJordan Gong <jordan.gong@protonmail.com>2021-02-26 20:17:30 +0800
committerJordan Gong <jordan.gong@protonmail.com>2021-02-26 20:17:30 +0800
commitaf4d28333b8204d6ab74738a8d3b6ae869cb9344 (patch)
treebb2bbc28256cb6290abcf62730a8537544ac9a3c /models
parent12f4b52df6f4ce8a4c4484146469774e4b158d30 (diff)
parent66d7b0e6f80aa6feb5e4e0af9af7ba78db641f9e (diff)
Merge branch 'python3.8' into python3.7
Diffstat (limited to 'models')
-rw-r--r--models/model.py3
-rw-r--r--models/rgb_part_net.py5
2 files changed, 5 insertions, 3 deletions
diff --git a/models/model.py b/models/model.py
index d179361..1e22483 100644
--- a/models/model.py
+++ b/models/model.py
@@ -311,7 +311,8 @@ class Model:
)
# Init models
- model_hp = self.hp.get('model', {})
+ model_hp: dict = self.hp.get('model', {}).copy()
+ model_hp.pop('triplet_margins', None)
self.rgb_pn = RGBPartNet(self.in_channels, self.in_size, **model_hp)
# Try to accelerate computation using CUDA or others
self.rgb_pn = self.rgb_pn.to(self.device)
diff --git a/models/rgb_part_net.py b/models/rgb_part_net.py
index fc1406c..4d7ba7f 100644
--- a/models/rgb_part_net.py
+++ b/models/rgb_part_net.py
@@ -76,8 +76,8 @@ class RGBPartNet(nn.Module):
def _disentangle(self, x_c1_t2, x_c2_t2=None):
n, t, c, h, w = x_c1_t2.size()
device = x_c1_t2.device
- x_c1_t1 = x_c1_t2[:, torch.randperm(t), :, :, :]
if self.training:
+ x_c1_t1 = x_c1_t2[:, torch.randperm(t), :, :, :]
((f_a_, f_c_, f_p_), losses) = self.ae(x_c1_t2, x_c1_t1, x_c2_t2)
# Decode features
x_c = self._decode_cano_feature(f_c_, n, t, device)
@@ -100,7 +100,8 @@ class RGBPartNet(nn.Module):
else: # evaluating
f_c_, f_p_ = self.ae(x_c1_t2)
x_c = self._decode_cano_feature(f_c_, n, t, device)
- x_p = self._decode_pose_feature(f_p_, n, t, c, h, w, device)
+ x_p_ = self._decode_pose_feature(f_p_, n, t, c, h, w, device)
+ x_p = x_p_.view(n, t, self.pn_in_channels, self.h // 4, self.w // 4)
return (x_c, x_p), None, None
def _decode_appr_feature(self, f_a_, n, t, device):