Upload train.py

train.py

@@ -166,11 +166,11 @@ def add_image_path(example):
     folder_name=example["folder_name"]
     #image_name = example['image_id'] + '.jpg'
     #image_path = os.path.join(r"D:\dump384_224x224_384\384", image_name)
-    image_path = [os.path.join(rf"~/dump_small/{folder_name[i]}", image_name[i]) for i in range(len(image_name))]
+    image_path = [os.path.join(rf"/home/user/dump_small/{folder_name[i]}", image_name[i]) for i in range(len(image_name))]
     example['image_path'] = image_path
     return example
 
-ds = dataset.map(add_image_path, batched=True, batch_size=1024)["train"]
+ds = dataset.map(add_image_path, batched=True, batch_size=8192)["train"]
 print(ds)
 
 ds = ds.train_test_split(test_size=0.02)
@@ -180,6 +180,8 @@ processed_dataset = ds.map(
     function=preprocess_fn,
     batched=True,
     fn_kwargs={"max_target_length": 128},
+    batch_size=8192,
+    num_proc=16,
     #remove_columns=ds['train'].column_names
 )