summaryrefslogtreecommitdiff
path: root/models/auto_encoder.py
diff options
context:
space:
mode:
authorJordan Gong <jordan.gong@protonmail.com>2021-04-10 22:38:23 +0800
committerJordan Gong <jordan.gong@protonmail.com>2021-04-10 22:38:23 +0800
commitf7df51490f9c3cc932672493e6c91595686df9ce (patch)
tree04975117a0f4a6b0d559a38091aba1db51a7a048 /models/auto_encoder.py
parent5e00cd7de1729db12329e793a4e84b6c7900a948 (diff)
parent20110729ab450c84d90965f5b8930236035f093a (diff)
Merge branch 'python3.8' into python3.7python3.7
Diffstat (limited to 'models/auto_encoder.py')
-rw-r--r--models/auto_encoder.py14
1 files changed, 1 insertions, 13 deletions
diff --git a/models/auto_encoder.py b/models/auto_encoder.py
index d11ec99..0c8bd5d 100644
--- a/models/auto_encoder.py
+++ b/models/auto_encoder.py
@@ -136,25 +136,13 @@ class AutoEncoder(nn.Module):
x_c1_t2_pred_ = self.decoder(f_a_c1_t1_, f_c_c1_t1_, f_p_c1_t2_)
x_c1_t2_pred = x_c1_t2_pred_.view(n, t, c, h, w)
- xrecon_loss = torch.stack([
- F.mse_loss(x_c1_t2[:, i], x_c1_t2_pred[:, i])
- for i in range(t)
- ]).sum()
-
f_c_c1_t1 = f_c_c1_t1_.view(f_size[1])
f_c_c2_t2 = f_c_c2_t2_.view(f_size[1])
- cano_cons_loss = torch.stack([
- F.mse_loss(f_c_c1_t1[:, i], f_c_c1_t2[:, i])
- + F.mse_loss(f_c_c1_t2[:, i], f_c_c2_t2[:, i])
- for i in range(t)
- ]).mean()
-
f_p_c2_t2 = f_p_c2_t2_.view(f_size[2])
- pose_sim_loss = F.mse_loss(f_p_c1_t2.mean(1), f_p_c2_t2.mean(1))
return (
(f_a_c1_t2, f_c_c1_t2, f_p_c1_t2),
- (xrecon_loss, cano_cons_loss, pose_sim_loss * 10)
+ (x_c1_t2_pred, (f_c_c1_t1, f_c_c2_t2), f_p_c2_t2)
)
else: # evaluating
return f_a_c1_t2, f_c_c1_t2, f_p_c1_t2