Updating transformer training code

Author: pythonlessons

Tutorials/05_sound_to_text/inferenceModel.py

@@ -36,7 +36,7 @@ def predict(self, data: np.ndarray):
         spectrogram = WavReader.get_spectrogram(wav_path, frame_length=configs.frame_length, frame_step=configs.frame_step, fft_length=configs.fft_length)
         # WavReader.plot_raw_audio(wav_path, label)
 
-        padded_spectrogram = np.pad(spectrogram, ((configs.max_spectrogram_length - spectrogram.shape[0], 0),(0,0)), mode="constant", constant_values=0)
+        padded_spectrogram = np.pad(spectrogram, ((0, configs.max_spectrogram_length - spectrogram.shape[0]),(0,0)), mode="constant", constant_values=0)
 
         # WavReader.plot_spectrogram(spectrogram, label)
 

Tutorials/09_translation_transformer/configs.py

@@ -14,13 +14,13 @@ def __init__(self):
         self.num_layers = 4
         self.d_model = 128
         self.num_heads = 8
-        self.dff = 128
+        self.dff = 512
         self.dropout_rate = 0.1
-        self.batch_size = 32
-        self.train_epochs = 20
+        self.batch_size = 16
+        self.train_epochs = 50
         # CustomSchedule parameters
         self.init_lr = 0.00001
         self.lr_after_warmup = 0.0005
         self.final_lr = 0.0001
         self.warmup_epochs = 2
-        self.decay_epochs = 9
+        self.decay_epochs = 18

Tutorials/09_translation_transformer/requirements.txt

@@ -0,0 +1 @@
+beautifulsoup4

Tutorials/09_translation_transformer/train.py

@@ -5,14 +5,13 @@
 except: pass
 
 from keras.callbacks import EarlyStopping, ModelCheckpoint, ReduceLROnPlateau, TensorBoard
-from mltu.tensorflow.callbacks import Model2onnx
+from mltu.tensorflow.callbacks import Model2onnx, WarmupCosineDecay
 
 from mltu.tensorflow.dataProvider import DataProvider
 from mltu.tokenizers import CustomTokenizer
 
 from mltu.tensorflow.transformer.utils import MaskedAccuracy, MaskedLoss
 from mltu.tensorflow.transformer.callbacks import EncDecSplitCallback
-from mltu.tensorflow.schedules import CustomSchedule
 
 from model import Transformer
 from configs import ModelConfigs
@@ -42,7 +41,7 @@ def read_files(path):
 es_training_data, en_training_data = zip(*train_dataset)
 es_validation_data, en_validation_data = zip(*val_dataset)
 
-# prepare portuguese tokenizer, this is the input language
+# prepare spanish tokenizer, this is the input language
 tokenizer = CustomTokenizer(char_level=True)
 tokenizer.fit_on_texts(es_training_data)
 tokenizer.save(configs.model_path + "/tokenizer.json")
@@ -99,17 +98,7 @@ def preprocess_inputs(data_batch, label_batch):
 
 transformer.summary()
 
-# Define learning rate schedule
-learning_rate = CustomSchedule(
-    steps_per_epoch=len(train_dataProvider),
-    init_lr=configs.init_lr,
-    lr_after_warmup=configs.lr_after_warmup,
-    final_lr=configs.final_lr,
-    warmup_epochs=configs.warmup_epochs,
-    decay_epochs=configs.decay_epochs,
-)
-
-optimizer = tf.keras.optimizers.Adam(learning_rate=learning_rate, beta_1=0.9, beta_2=0.98, epsilon=1e-9)
+optimizer = tf.keras.optimizers.Adam(learning_rate=configs.init_lr, beta_1=0.9, beta_2=0.98, epsilon=1e-9)
 
 # Compile the model
 transformer.compile(
@@ -120,6 +109,13 @@ def preprocess_inputs(data_batch, label_batch):
     )
 
 # Define callbacks
+warmupCosineDecay = WarmupCosineDecay(
+    lr_after_warmup=configs.lr_after_warmup,
+    final_lr=configs.final_lr,
+    warmup_epochs=configs.warmup_epochs,
+    decay_epochs=configs.decay_epochs,
+    initial_lr=configs.init_lr,
+    )
 earlystopper = EarlyStopping(monitor="val_masked_accuracy", patience=5, verbose=1, mode="max")
 checkpoint = ModelCheckpoint(f"{configs.model_path}/model.h5", monitor="val_masked_accuracy", verbose=1, save_best_only=True, mode="max", save_weights_only=False)
 tb_callback = TensorBoard(f"{configs.model_path}/logs")

mltu/tensorflow/callbacks.py

@@ -101,4 +101,48 @@ def __init__(self, log_path: str, log_file: str="logs.log", logLevel=logging.INF
     def on_epoch_end(self, epoch: int, logs: dict=None):
         epoch_message = f"Epoch {epoch}; "
         logs_message = "; ".join([f"{key}: {value}" for key, value in logs.items()])
-        self.logger.info(epoch_message + logs_message)
+        self.logger.info(epoch_message + logs_message)
+
+
+class WarmupCosineDecay(Callback):
+    """ Cosine decay learning rate scheduler with warmup
+
+    Args:
+        lr_after_warmup (float): Learning rate after warmup
+        final_lr (float): Final learning rate
+        warmup_epochs (int): Number of warmup epochs
+        decay_epochs (int): Number of decay epochs
+        initial_lr (float, optional): Initial learning rate. Defaults to 0.0.
+        verbose (bool, optional): Whether to print learning rate. Defaults to False.
+    """
+    def __init__(
+            self, 
+            lr_after_warmup: float, 
+            final_lr: float, 
+            warmup_epochs: int, 
+            decay_epochs: int, 
+            initial_lr: float=0.0, 
+            verbose=False
+        ) -> None:
+        super(WarmupCosineDecay, self).__init__()
+        self.lr_after_warmup = lr_after_warmup
+        self.final_lr = final_lr
+        self.warmup_epochs = warmup_epochs
+        self.decay_epochs = decay_epochs
+        self.initial_lr = initial_lr
+        self.verbose = verbose
+
+    def on_epoch_begin(self, epoch: int, logs: dict=None):
+        """ Adjust learning rate at the beginning of each epoch """
+        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:
+            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}")

mltu/tensorflow/schedules.py

@@ -204,6 +204,14 @@ def __init__(self, num_layers: int, d_model: int, num_heads: int, dff: int, voca
             for _ in range(num_layers)]
         self.dropout = tf.keras.layers.Dropout(dropout_rate)
 
+    def get_config(self):
+        config = super().get_config()
+        config.update({
+            'd_model': self.d_model,
+            'num_layers': self.num_layers,
+        })
+        return config
+
     def call(self, x: tf.Tensor) -> tf.Tensor:
         """
         The call function that performs the forward pass of the layer.
@@ -323,6 +331,14 @@ def __init__(self, num_layers: int, d_model: int, num_heads: int, dff: int, voca
 
         self.last_attn_scores = None
 
+    def get_config(self):
+        config = super().get_config()
+        config.update({
+            'd_model': self.d_model,
+            'num_layers': self.num_layers,
+        })
+        return config
+
     def call(self, x: tf.Tensor, context: tf.Tensor) -> tf.Tensor:
         """
         The call function that performs the forward pass of the layer.

mltu/tensorflow/transformer/layers.py

@@ -181,7 +181,8 @@ def save(self, path: str, type: str="json"):
         """
         serialised_dict = self.dict()
         if type == "json":
-            os.makedirs(os.path.dirname(path), exist_ok=True)
+            if os.path.dirname(path):
+                os.makedirs(os.path.dirname(path), exist_ok=True)
             with open(path, "w") as f:
                 json.dump(serialised_dict, f)
 

mltu/tokenizers.py

@@ -159,7 +159,7 @@ def __init__(
         self.padding_value = padding_value
 
     def __call__(self, spectrogram: np.ndarray, label: np.ndarray):
-        padded_spectrogram = np.pad(spectrogram, ((self.max_spectrogram_length - spectrogram.shape[0], 0),(0,0)), mode="constant", constant_values=self.padding_value)
+        padded_spectrogram = np.pad(spectrogram, (0, (self.max_spectrogram_length - spectrogram.shape[0]),(0,0)), mode="constant", constant_values=self.padding_value)
 
         return padded_spectrogram, label