From 06e9a53673fb193f8287d9d9b95463a5d1b044bb Mon Sep 17 00:00:00 2001 From: Jordan Gong Date: Fri, 5 Mar 2021 20:08:22 +0800 Subject: Calculate losses outside modules --- models/auto_encoder.py | 15 +++------------ 1 file changed, 3 insertions(+), 12 deletions(-) (limited to 'models/auto_encoder.py') diff --git a/models/auto_encoder.py b/models/auto_encoder.py index e6a3e60..0694ff1 100644 --- a/models/auto_encoder.py +++ b/models/auto_encoder.py @@ -151,27 +151,18 @@ 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(n, t, -1) f_c_c1_t2 = f_c_c1_t2_.view(n, t, -1) f_c_c2_t2 = f_c_c2_t2_.view(n, t, -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_c1_t2 = f_p_c1_t2_.view(n, t, -1) f_p_c2_t2 = f_p_c2_t2_.view(n, t, -1) - 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_), - torch.stack((xrecon_loss, cano_cons_loss, pose_sim_loss * 10)) + (x_c1_t2_pred, + (f_c_c1_t1, f_c_c1_t2, f_c_c2_t2), + (f_p_c1_t2, f_p_c2_t2)) ) else: # evaluating return f_c_c1_t2_, f_p_c1_t2_ -- cgit v1.2.3