Change MDD loss + Add MCD multiple steps encoder

Author: antoinedemathelin

adapt/feature_based/_mcd.py

@@ -141,6 +141,26 @@ def train_step(self, data):
             Xs = Xs[0]
             ys = ys[0]
 
+            
+            for _ in range(4):
+                with tf.GradientTape() as enc_tape:
+                    Xt_enc = self.encoder_(Xt, training=True)
+                    yt_pred = self.task_(Xt_enc, training=True)
+                    yt_disc = self.discriminator_(Xt_enc, training=True)
+
+                    # Reshape
+                    yt_pred = tf.reshape(yt_pred, tf.shape(ys))
+                    yt_disc = tf.reshape(yt_disc, tf.shape(ys))
+
+                    discrepancy = tf.reduce_mean(tf.abs(yt_pred - yt_disc))
+                    enc_loss = discrepancy
+                    enc_loss += sum(self.encoder_.losses)
+                    
+                # Compute gradients
+                trainable_vars_enc = self.encoder_.trainable_variables
+                gradients_enc = enc_tape.gradient(enc_loss, trainable_vars_enc)
+                self.optimizer.apply_gradients(zip(gradients_enc, trainable_vars_enc))
+            
             # loss
             with tf.GradientTape() as task_tape, tf.GradientTape() as enc_tape, tf.GradientTape() as disc_tape:
                 # Forward pass
@@ -174,7 +194,6 @@ def train_step(self, data):
                 disc_loss += sum(self.discriminator_.losses)
                 enc_loss += sum(self.encoder_.losses)
 
-
             # Compute gradients
             trainable_vars_task = self.task_.trainable_variables
             trainable_vars_enc = self.encoder_.trainable_variables

adapt/feature_based/_mdd.py

@@ -25,7 +25,7 @@ class MDD(BaseAdaptDeep):
     
     Parameters
     ----------
-    lambda_ : float (default=1.)
+    lambda_ : float (default=0.1)
         Trade-off parameter
     
     gamma : float (default=4.)
@@ -71,7 +71,7 @@ def __init__(self,
                  task=None,
                  Xt=None,
                  yt=None,
-                 lambda_=1.,
+                 lambda_=0.1,
                  gamma=4.,
                  copy=True,
                  verbose=1,
@@ -124,8 +124,8 @@ def train_step(self, data):
                              tf.shape(ys_pred)[1])
                 argmax_tgt = tf.one_hot(tf.math.argmax(yt_pred, -1),
                              tf.shape(yt_pred)[1])
-                disc_loss_src = self.task_loss_(argmax_src, ys_disc)
-                disc_loss_tgt = self.task_loss_(argmax_tgt, yt_disc)
+                disc_loss_src = -tf.math.log(tf.reduce_sum(argmax_src * ys_disc, 1) + EPS)
+                disc_loss_tgt = tf.math.log(1. - tf.reduce_sum(argmax_tgt * yt_disc, 1) + EPS)
             else:
                 disc_loss_src = self.task_loss_(ys_pred, ys_disc)
                 disc_loss_tgt = self.task_loss_(yt_pred, yt_disc)
@@ -168,6 +168,7 @@ def train_step(self, data):
         logs = {m.name: m.result() for m in self.metrics}
         # disc_metrics = self._get_disc_metrics(ys_disc, yt_disc)
         logs.update({"disc_loss": disc_loss})
+        logs.update({"disc_src": disc_loss_src, "disc_tgt": disc_loss_tgt})
         return logs