add cache to mem

main.py

@@ -11,6 +11,7 @@ import os, glob
 from callbacks import BreakEachEpoch
 import subprocess
 from multiprocessing import Process
+import shutil
 
 logging.set_verbosity_info()
 
@@ -71,7 +72,7 @@ def prepare_dataset(batch, processor):
     return batch
 
 
-def load_prepared_dataset(path, processor, cache_file_filter_name, cache_file_map_name, num_proc=6):
+def load_prepared_dataset(path, processor, cache_file_filter_name, cache_file_map_name, num_proc=5):
     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(
@@ -88,7 +89,7 @@ def load_prepared_dataset(path, processor, cache_file_filter_name, cache_file_ma
     # 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,
+            shutil.move(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,
@@ -128,7 +129,7 @@ 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))
+    _test_dataset_shard_idx = min(round(_train_dataset_shard_idx / (num_train_shards / num_test_shards)), num_test_shards - 1)
     _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
@@ -171,7 +172,7 @@ if __name__ == "__main__":
     train_dataset_root_folder = '/content/drive/MyDrive/audio_dataset/train_dataset'
     test_dataset_root_folder = '/content/drive/MyDrive/audio_dataset/test_dataset'
 
-    cache_processing_dataset_folder = './data-bin/cache/'
+    cache_processing_dataset_folder = '/dev/shm/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'))