Fixed learning rate logging from warmupCosineDecay callback

Author: pythonlessons

Tutorials/09_translation_transformer/train.py

@@ -72,15 +72,15 @@ def preprocess_inputs(data_batch, label_batch):
     train_dataset, 
     batch_size=configs.batch_size, 
     batch_postprocessors=[preprocess_inputs],
-    use_cache=True
+    use_cache=True,
     )
 
 # Create Validation Data Provider
 val_dataProvider = DataProvider(
     val_dataset, 
     batch_size=configs.batch_size, 
     batch_postprocessors=[preprocess_inputs],
-    use_cache=True
+    use_cache=True,
     )
 
 # Create TensorFlow Transformer Model
@@ -129,6 +129,7 @@ def preprocess_inputs(data_batch, label_batch):
     validation_data=val_dataProvider, 
     epochs=configs.train_epochs,
     callbacks=[
+        warmupCosineDecay,
         checkpoint, 
         tb_callback, 
         reduceLROnPlat,

mltu/tensorflow/callbacks.py

@@ -134,15 +134,25 @@ def __init__(
 
     def on_epoch_begin(self, epoch: int, logs: dict=None):
         """ Adjust learning rate at the beginning of each epoch """
+
+        if epoch >= self.warmup_epochs + self.decay_epochs:
+            return logs
+
         if epoch < self.warmup_epochs:
             lr = self.initial_lr + (self.lr_after_warmup - self.initial_lr) * (epoch + 1) / self.warmup_epochs
-        elif epoch < self.warmup_epochs + self.decay_epochs:
+        else:
             progress = (epoch - self.warmup_epochs) / self.decay_epochs
             lr = self.final_lr + 0.5 * (self.lr_after_warmup - self.final_lr) * (1 + tf.cos(tf.constant(progress) * 3.14159))
-        else:
-            return None # No change to learning rate
 
         tf.keras.backend.set_value(self.model.optimizer.lr, lr)
 
         if self.verbose:
-            print(f"Epoch {epoch + 1} - Learning Rate: {lr}")
+            print(f"Epoch {epoch + 1} - Learning Rate: {lr}")
+    
+    def on_epoch_end(self, epoch: int, logs: dict=None):
+        logs = logs or {}
+        
+        # Log the learning rate value
+        logs["lr"] = self.model.optimizer.lr
+        
+        return logs

mltu/tensorflow/transformer/utils.py

@@ -76,29 +76,4 @@ def result(self) -> tf.Tensor:
         Returns:
             tf.Tensor: Masked accuracy.
         """
-        return self.total / self.count
-
-
-# def masked_accuracy(y_true: tf.Tensor, y_pred: tf.Tensor):
-#     """ Calculate masked accuracy.
-
-#     Args:
-#         y_true (tf.Tensor): True labels.
-#         y_pred (tf.Tensor): Predicted labels.
-
-#     Returns:
-#         tf.Tensor: Masked accuracy.
-#     """
-#     pred = tf.argmax(y_pred, axis=2)
-#     label = tf.cast(y_true, pred.dtype)
-#     match = label == pred
-
-#     mask = label != 0
-
-#     match = match & mask
-
-#     match = tf.cast(match, dtype=tf.float32)
-#     mask = tf.cast(mask, dtype=tf.float32)
-#     accuracy = tf.reduce_sum(match) / tf.reduce_sum(mask)
-
-#     return accuracy
+        return self.total / self.count