add prefetch

main.py

@@ -10,6 +10,7 @@ import json
 import os, glob
 from callbacks import BreakEachEpoch
 import subprocess
+from multiprocessing import Process
 
 logging.set_verbosity_info()
 
@@ -70,8 +71,34 @@ def prepare_dataset(batch, processor):
     return batch
 
 
-def load_prepared_dataset(path, processor, cache_file_filter_name, cache_file_map_name, num_proc=8):
+def load_prepared_dataset(path, processor, cache_file_filter_name, cache_file_map_name, num_proc=6):
     dataset = load_from_disk(path)
+    list_cache_prefetch_files = glob.glob(
+        cache_file_map_name.replace(cache_processing_dataset_folder, cache_processing_dataset_folder_prefetch).replace(
+            '.arrow', '*'))
+
+    # Do not re-compute what already in cache folder
+    if cache_file_map_name.startswith(cache_processing_dataset_folder_prefetch):
+        if len(glob.glob(cache_file_map_name.replace(cache_processing_dataset_folder_prefetch,
+                                                 cache_processing_dataset_folder).replace('.arrow', '*'))) > 0:
+            return
+        if len(list_cache_prefetch_files) > 0:
+            return
+
+    # check cache file
+    if len(glob.glob(cache_file_map_name.replace('.arrow', '*'))) == 0 and len(list_cache_prefetch_files) > 0:
+        for item_file in list_cache_prefetch_files:
+            os.rename(item_file, item_file.replace(cache_processing_dataset_folder_prefetch,
+                                                   cache_processing_dataset_folder))
+    if len(glob.glob(cache_file_map_name.replace('.arrow', '*'))) > 0:
+        return dataset.map(prepare_dataset,
+                           remove_columns=dataset.column_names,
+                           batch_size=32,
+                           num_proc=num_proc,
+                           batched=True,
+                           fn_kwargs={"processor": processor},
+                           cache_file_name=cache_file_map_name)
+
     dataset = dataset.filter(lambda example: len(example['speech']) < 160000,
                              batch_size=32,
                              num_proc=num_proc,
@@ -83,6 +110,7 @@ def load_prepared_dataset(path, processor, cache_file_filter_name, cache_file_ma
                                     batched=True,
                                     fn_kwargs={"processor": processor},
                                     cache_file_name=cache_file_map_name)
+    processed_dataset.cleanup_cache_files()
     return processed_dataset
 
 
@@ -95,6 +123,44 @@ def commit_checkpoint():
     for command in submit_commands:
         print(subprocess.run(command.split(), stdout=subprocess.PIPE).stdout.decode('utf-8'))
 
+
+def get_train_test_shard_id(epoch_count):
+    # loop over training shards
+    _train_dataset_shard_idx = epoch_count % num_train_shards
+    # Get test shard depend on train shard id
+    _test_dataset_shard_idx = round(_train_dataset_shard_idx / (num_train_shards / num_test_shards))
+    _num_test_sub_shard = 8  # Split test shard into subset. Default is 8
+    _idx_sub_shard = _train_dataset_shard_idx % _num_test_sub_shard  # loop over test shard subset
+    return _train_dataset_shard_idx, _test_dataset_shard_idx, _num_test_sub_shard, _idx_sub_shard
+
+
+def process_prefetch_epoch(epoch_count):
+    train_shard_idx, test_shard_idx, _, _ = get_train_test_shard_id(epoch_count)
+    load_prepared_dataset(os.path.join(train_dataset_root_folder,
+                                       'shard_{}'.format(train_shard_idx)),
+                          w2v_ctc_processor,
+                          cache_file_filter_name=os.path.join(cache_processing_dataset_folder_prefetch,
+                                                              'train',
+                                                              'cache-train-filter-shard-{}.arrow'.format(
+                                                                  train_shard_idx)),
+                          cache_file_map_name=os.path.join(cache_processing_dataset_folder_prefetch,
+                                                           'train',
+                                                           'cache-train-map-shard-{}.arrow'.format(
+                                                               train_shard_idx)),
+                          )
+    load_prepared_dataset(os.path.join(test_dataset_root_folder,
+                                       'shard_{}'.format(test_shard_idx)),
+                          w2v_ctc_processor,
+                          cache_file_filter_name=os.path.join(cache_processing_dataset_folder_prefetch,
+                                                              'test',
+                                                              'cache-test-filter-shard-{}.arrow'.format(
+                                                                  test_shard_idx)),
+                          cache_file_map_name=os.path.join(cache_processing_dataset_folder_prefetch, 'test',
+                                                           'cache-test-map-shard-{}.arrow'.format(
+                                                               test_shard_idx))
+                          )
+
+
 if __name__ == "__main__":
 
     checkpoint_path = "./model-bin/finetune/base/"
@@ -106,9 +172,13 @@ if __name__ == "__main__":
     test_dataset_root_folder = '/content/drive/MyDrive/audio_dataset/test_dataset'
 
     cache_processing_dataset_folder = './data-bin/cache/'
+    cache_processing_dataset_folder_prefetch = './data-bin/cache_prefetch/'
     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'))
+    if not os.path.exists(os.path.join(cache_processing_dataset_folder_prefetch, 'train')):
+        os.makedirs(os.path.join(cache_processing_dataset_folder_prefetch, 'train'))
+        os.makedirs(os.path.join(cache_processing_dataset_folder_prefetch, '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
@@ -121,7 +191,7 @@ if __name__ == "__main__":
         per_device_eval_batch_size=32,
         gradient_accumulation_steps=2,
         num_train_epochs=num_epochs,  # each epoch per shard data
-        logging_steps=1,
+        logging_steps=5,
         learning_rate=1e-5,
         weight_decay=0.005,
         warmup_steps=1000,
@@ -150,13 +220,23 @@ if __name__ == "__main__":
     w2v_ctc_model, w2v_ctc_processor = load_pretrained_model()
     data_collator = DataCollatorCTCWithPadding(processor=w2v_ctc_processor, padding=True)
 
+    prefetch_process = []
+
     for epoch_idx in range(last_epoch_idx, num_epochs):
-        # loop over training shards
-        train_dataset_shard_idx = epoch_idx % num_train_shards
-        # Get test shard depend on train shard id
-        test_dataset_shard_idx = round(train_dataset_shard_idx / (num_train_shards / num_test_shards))
-        num_test_sub_shard = 8  # Split test shard into subset. Default is 8
-        idx_sub_shard = train_dataset_shard_idx % num_test_sub_shard  # loop over test shard subset
+        # # loop over training shards
+        # train_dataset_shard_idx = epoch_idx % num_train_shards
+        # # Get test shard depend on train shard id
+        # test_dataset_shard_idx = round(train_dataset_shard_idx / (num_train_shards / num_test_shards))
+        # num_test_sub_shard = 8  # Split test shard into subset. Default is 8
+        # idx_sub_shard = train_dataset_shard_idx % num_test_sub_shard  # loop over test shard subset
+
+        train_dataset_shard_idx, test_dataset_shard_idx, num_test_sub_shard, idx_sub_shard = get_train_test_shard_id(
+            epoch_idx)
+
+        # waiting for all prefetch process done
+        for process_instance in prefetch_process:
+            process_instance.join()
+        prefetch_process.clear()
 
         # load train shard
         train_dataset = load_prepared_dataset(os.path.join(train_dataset_root_folder,
@@ -170,7 +250,7 @@ if __name__ == "__main__":
                                                                                'train',
                                                                                'cache-train-map-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)),
@@ -184,6 +264,12 @@ if __name__ == "__main__":
                                                                                   test_dataset_shard_idx))
                                              )
         test_dataset = test_dataset.shard(num_test_sub_shard, idx_sub_shard)
+
+        # Prefetch_dataset
+        prefetch_process.append(Process(target=process_prefetch_epoch, args=(epoch_idx + 1,)))
+        for process_instance in prefetch_process:
+            process_instance.start()
+
         # Init trainer
         if trainer is None:
             trainer = Trainer(
@@ -216,5 +302,5 @@ if __name__ == "__main__":
         test_dataset.cleanup_cache_files()
         train_dataset.cleanup_cache_files()
 
-        if epoch_idx % 10 == 0:
-          commit_checkpoint()
+        if epoch_idx % 5 == 0:
+            commit_checkpoint()