Add cross entropy lossdisentangling_only

1 files changed, 7 insertions, 2 deletions

diff --git a/models/auto_encoder.py b/models/auto_encoder.py
index 91071dd..9bfb365 100644
--- a/models/auto_encoder.py
+++ b/models/auto_encoder.py
@@ -110,7 +110,7 @@ class Decoder(nn.Module):
         x = torch.cat((f_appearance, f_canonical, f_pose), dim=1)
         x = self.fc(x)
         x = x.view(-1, self.feature_channels * 8, self.h_0, self.w_0)
-        x = F.relu(x, inplace=True)
+        x = F.leaky_relu(x, 0.2, inplace=True)
         x = self.trans_conv1(x)
         x = self.trans_conv2(x)
         x = self.trans_conv3(x)
@@ -122,6 +122,7 @@ class Decoder(nn.Module):
 class AutoEncoder(nn.Module):
     def __init__(
             self,
+            num_class: int,
             channels: int = 3,
             frame_size: tuple[int, int] = (64, 48),
             feature_channels: int = 64,
@@ -132,8 +133,9 @@ class AutoEncoder(nn.Module):
                                feature_channels, embedding_dims)
         self.decoder = Decoder(embedding_dims, feature_channels,
                                self.encoder.feature_size, channels)
+        self.classifier = BasicLinear(embedding_dims[1], num_class)
 
-    def forward(self, x_c1_t2, x_c1_t1=None, x_c2_t2=None):
+    def forward(self, x_c1_t2, x_c1_t1=None, x_c2_t2=None, y=None):
         n, t, c, h, w = x_c1_t2.size()
         # x_c1_t2 is the frame for later module
         x_c1_t2_ = x_c1_t2.view(n * t, c, h, w)
@@ -160,6 +162,9 @@ class AutoEncoder(nn.Module):
             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, :])
+                + F.cross_entropy(self.classifier(
+                    F.leaky_relu(f_c_c1_t2[:, i, :], 0.2)
+                ), y)
                 for i in range(t)
             ]).mean()