add config for training multi epochs

callbacks.py

@@ -0,0 +1,12 @@
+from transformers import TrainerCallback, TrainingArguments, TrainerState, TrainerControl, logging
+
+
+class BreakEachEpoch(TrainerCallback):
+    """
+    A :class:`~transformers.TrainerCallback` that handles the default flow of the training loop for logs, evaluation
+    and checkpoints.
+    """
+    def on_epoch_end(self, args: TrainingArguments, state: TrainerState, control: TrainerControl, **kwargs):
+        control.should_training_stop = True
+        logging.get_logger().info("Break each epoch for reload new shard dataset")
+        return control

main.py

@@ -1,12 +1,14 @@
-from transformers import Wav2Vec2ForCTC, Wav2Vec2Processor, Wav2Vec2CTCTokenizer, Wav2Vec2FeatureExtractor
+from transformers import Wav2Vec2ForCTC, Wav2Vec2Processor, Wav2Vec2CTCTokenizer, Wav2Vec2FeatureExtractor, \
+    TrainerCallback
 from datasets import load_from_disk
 from data_handler import DataCollatorCTCWithPadding
 from transformers import TrainingArguments
 from transformers import Trainer, logging
 from metric_utils import compute_metrics_fn
 from transformers.trainer_utils import get_last_checkpoint
-import json, random
+import json
 import os, glob
+from callbacks import BreakEachEpoch
 
 logging.set_verbosity_info()
 
@@ -68,8 +70,8 @@ def load_prepared_dataset(path, processor, cache_file_name):
     dataset = load_from_disk(path)
     processed_dataset = dataset.map(prepare_dataset,
                                     remove_columns=dataset.column_names,
-                                    batch_size=8,
-                                    num_proc=8,
+                                    batch_size=32,
+                                    num_proc=4,
                                     batched=True,
                                     fn_kwargs={"processor": processor},
                                     cache_file_name=cache_file_name)
@@ -90,8 +92,9 @@ if __name__ == "__main__":
     test_dataset_root_folder = '/content/drive/MyDrive/audio_dataset/test_dataset'
 
     cache_processing_dataset_folder = './data-bin/cache/'
-    if not os.path.exists(cache_processing_dataset_folder):
-        os.makedirs(cache_processing_dataset_folder)
+    if not os.path.exists(os.path.join(cache_processing_dataset_folder, 'train')):
+        os.makedirs(os.path.join(cache_processing_dataset_folder, 'train'))
+        os.makedirs(os.path.join(cache_processing_dataset_folder, 'test'))
     num_train_shards = len(glob.glob(os.path.join(train_dataset_root_folder, 'shard_*')))
     num_test_shards = len(glob.glob(os.path.join(test_dataset_root_folder, 'shard_*')))
     num_epochs = 5000
@@ -100,20 +103,21 @@ if __name__ == "__main__":
         output_dir=checkpoint_path,
         # fp16=True,
         group_by_length=True,
-        per_device_train_batch_size=16,
-        per_device_eval_batch_size=16,
+        per_device_train_batch_size=4,
+        per_device_eval_batch_size=4,
         gradient_accumulation_steps=8,
-        num_train_epochs=1,  # each epoch per shard data
+        num_train_epochs=num_epochs,  # each epoch per shard data
         logging_steps=1,
         learning_rate=1e-4,
         weight_decay=0.005,
-        warmup_steps=5000,
+        warmup_steps=1000,
         save_total_limit=2,
         ignore_data_skip=True,
         logging_dir=os.path.join(checkpoint_path, 'log'),
         metric_for_best_model='wer',
         save_strategy="epoch",
         evaluation_strategy="epoch",
+        greater_is_better=False,
         # save_steps=5,
         # eval_steps=5,
     )
@@ -143,19 +147,19 @@ if __name__ == "__main__":
         train_dataset = load_prepared_dataset(os.path.join(train_dataset_root_folder,
                                                            'shard_{}'.format(train_dataset_shard_idx)),
                                               w2v_ctc_processor,
-                                              cache_file_name=os.path.join(cache_processing_dataset_folder,
+                                              cache_file_name=os.path.join(cache_processing_dataset_folder, 'train',
                                                                            'cache-train-shard-{}.arrow'.format(
                                                                                train_dataset_shard_idx))
-                                              )  # .shard(1000, 0)  # Remove shard split when train
+                                              ).shard(1000, 0)  # Remove shard split when train
         # load test shard subset
         test_dataset = load_prepared_dataset(os.path.join(test_dataset_root_folder,
                                                           'shard_{}'.format(test_dataset_shard_idx)),
                                              w2v_ctc_processor,
-                                             cache_file_name=os.path.join(cache_processing_dataset_folder,
+                                             cache_file_name=os.path.join(cache_processing_dataset_folder, 'test',
                                                                           'cache-test-shard-{}.arrow'.format(
                                                                               test_dataset_shard_idx))
-                                             ).shard(num_test_sub_shard, idx_sub_shard)
-
+                                             )
+        test_dataset = test_dataset.shard(num_test_sub_shard, idx_sub_shard)
         # Init trainer
         trainer = Trainer(
             model=w2v_ctc_model,
@@ -164,13 +168,16 @@ if __name__ == "__main__":
             compute_metrics=compute_metrics_fn(w2v_ctc_processor),
             train_dataset=train_dataset,
             eval_dataset=test_dataset,
-            tokenizer=w2v_ctc_processor.feature_extractor
+            tokenizer=w2v_ctc_processor.feature_extractor,
+            callbacks=[BreakEachEpoch()]  # Manual break end of epoch because each epoch loop over a shard
         )
-        # Manual add num_train_epochs because each epoch loop over a shard
-        training_args.num_train_epochs = epoch_idx + 1
 
-        logging.get_logger().info('Train shard idx: {}'.format(train_dataset_shard_idx))
-        logging.get_logger().info('Valid shard idx: {} sub_shard: {}'.format(test_dataset_shard_idx, idx_sub_shard))
+        # training_args.num_train_epochs = epoch_idx + 1
+
+        logging.get_logger().info('Train epoch {}'.format(training_args.num_train_epochs))
+        logging.get_logger().info('Train shard idx: {} / {}'.format(train_dataset_shard_idx + 1, num_train_shards))
+        logging.get_logger().info(
+            'Valid shard idx: {} / {} sub_shard: {}'.format(test_dataset_shard_idx + 1, num_test_shards, idx_sub_shard))
 
         if last_checkpoint_path is not None:
             # start train from a checkpoint if exist
@@ -181,5 +188,5 @@ if __name__ == "__main__":
         last_checkpoint_path = get_last_checkpoint(checkpoint_path)
 
         # Clear cache file to free disk
-        # test_dataset.cleanup_cache_files()
-        # train_dataset.cleanup_cache_files()
+        test_dataset.cleanup_cache_files()
+        train_dataset.cleanup_cache_files()