Added loss function from 10.7.5 to Seq2Seq

Author: nickmcgreivy

d2l/jax.py

@@ -1199,6 +1199,17 @@ def validation_step(self, params, batch, state):
     def configure_optimizers(self):
         # Adam optimizer is used here
         return optax.adam(learning_rate=self.lr)
+    
+    @partial(jax.jit, static_argnums=(0, 5))
+    def loss(self, params, X, Y, state, averaged=False):
+        Y_hat = state.apply_fn({'params': params}, *X,
+                            rngs={'dropout': state.dropout_rng})
+        Y_hat = Y_hat.reshape((-1, Y_hat.shape[-1]))
+        Y = Y.reshape((-1,))
+        fn = optax.softmax_cross_entropy_with_integer_labels
+        l = fn(Y_hat, Y)
+        mask = (Y.reshape(-1) != self.tgt_pad).astype(jnp.float32)
+        return (l * mask).sum() / mask.sum(), {}
 
 def bleu(pred_seq, label_seq, k):
     """Compute the BLEU.

d2l/mxnet.py

@@ -1025,6 +1025,11 @@ def configure_optimizers(self):
         # Adam optimizer is used here
         return gluon.Trainer(self.parameters(), 'adam',
                              {'learning_rate': self.lr})
+    
+    def loss(self, Y_hat, Y):
+        l = super(Seq2Seq, self).loss(Y_hat, Y, averaged=False)
+        mask = (Y.reshape(-1) != self.tgt_pad).astype(np.float32)
+        return (l * mask).sum() / mask.sum()
 
 def bleu(pred_seq, label_seq, k):
     """Compute the BLEU.

d2l/tensorflow.py

@@ -979,6 +979,11 @@ def configure_optimizers(self):
         # Adam optimizer is used here
         return tf.keras.optimizers.Adam(learning_rate=self.lr)
 
+    def loss(self, Y_hat, Y):
+        l = super(Seq2Seq, self).loss(Y_hat, Y, averaged=False)
+        mask = tf.cast(tf.reshape(Y, -1) != self.tgt_pad, tf.float32)
+        return tf.reduce_sum(l * mask) / tf.reduce_sum(mask)
+
 def bleu(pred_seq, label_seq, k):
     """Compute the BLEU.
 

d2l/torch.py

@@ -1026,6 +1026,11 @@ def configure_optimizers(self):
         # Adam optimizer is used here
         return torch.optim.Adam(self.parameters(), lr=self.lr)
 
+    def loss(self, Y_hat, Y):
+        l = super(Seq2Seq, self).loss(Y_hat, Y, averaged=False)
+        mask = (Y.reshape(-1) != self.tgt_pad).type(torch.float32)
+        return (l * mask).sum() / mask.sum()
+
 def bleu(pred_seq, label_seq, k):
     """Compute the BLEU.