filter wav 10s and new pretrained model

main.py

@@ -45,6 +45,7 @@ def load_pretrained_model(checkpoint_path=None):
         )
         # model.freeze_feature_extractor()
 
+    # model = Wav2Vec2ForCTC(model.config)
     model_total_params = sum(p.numel() for p in model.parameters())
     model_total_params_trainable = sum(p.numel() for p in model.parameters() if p.requires_grad)
     print(model)
@@ -68,15 +69,19 @@ def prepare_dataset(batch, processor):
     return batch
 
 
-def load_prepared_dataset(path, processor, cache_file_name):
+def load_prepared_dataset(path, processor, cache_file_filter_name, cache_file_map_name, num_proc=8):
     dataset = load_from_disk(path)
+    dataset = dataset.filter(lambda example: len(example['speech']) < 160000,
+                             batch_size=32,
+                             num_proc=num_proc,
+                             cache_file_name=cache_file_filter_name)
     processed_dataset = dataset.map(prepare_dataset,
                                     remove_columns=dataset.column_names,
                                     batch_size=32,
-                                    num_proc=4,
+                                    num_proc=num_proc,
                                     batched=True,
                                     fn_kwargs={"processor": processor},
-                                    cache_file_name=cache_file_name)
+                                    cache_file_name=cache_file_map_name)
     return processed_dataset
 
 
@@ -105,9 +110,9 @@ if __name__ == "__main__":
         output_dir=checkpoint_path,
         fp16=True,
         group_by_length=True,
-        per_device_train_batch_size=4,
-        per_device_eval_batch_size=4,
-        gradient_accumulation_steps=8,
+        per_device_train_batch_size=32,
+        per_device_eval_batch_size=32,
+        gradient_accumulation_steps=2,
         num_train_epochs=num_epochs,  # each epoch per shard data
         logging_steps=1,
         learning_rate=1e-4,
@@ -150,17 +155,26 @@ 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, 'train',
-                                                                           'cache-train-shard-{}.arrow'.format(
-                                                                               train_dataset_shard_idx))
-                                              )  # .shard(1000, 0)  # Remove shard split when train
+                                              cache_file_filter_name=os.path.join(cache_processing_dataset_folder,
+                                                                                  'train',
+                                                                                  'cache-train-filter-shard-{}.arrow'.format(
+                                                                                      train_dataset_shard_idx)),
+                                              cache_file_map_name=os.path.join(cache_processing_dataset_folder,
+                                                                               'train',
+                                                                               'cache-train-map-shard-{}.arrow'.format(
+                                                                                   train_dataset_shard_idx)),
+                                              ) #.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, 'test',
-                                                                          'cache-test-shard-{}.arrow'.format(
-                                                                              test_dataset_shard_idx))
+                                             cache_file_filter_name=os.path.join(cache_processing_dataset_folder,
+                                                                                 'test',
+                                                                                 'cache-test-filter-shard-{}.arrow'.format(
+                                                                                     test_dataset_shard_idx)),
+                                             cache_file_map_name=os.path.join(cache_processing_dataset_folder, 'test',
+                                                                              'cache-test-map-shard-{}.arrow'.format(
+                                                                                  test_dataset_shard_idx))
                                              )
         test_dataset = test_dataset.shard(num_test_sub_shard, idx_sub_shard)
         # Init trainer

model-bin/pretrained/base/pytorch_model.bin

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:8b36355988e4d1f94d070ef677ab4d304bce440af0c3dd7bd1c98e295e907f09
+oid sha256:b8fc5e67c00d407cd160a238034677db5670cbc77fe766c53d1042478509574d
 size 380261837