Update train.py

train.py

@@ -115,34 +115,35 @@ def process_shard_worker(args):
 
     return result
 
-mp.set_start_method("spawn", force=True)
+if __name__ == '__main__':
+    mp.set_start_method("spawn", force=True)
 
-input_dataset = "none-yet/anime-captions"
-output_dataset = "nroggendorff/anime-captions"
-model_name = "datalab-to/chandra"
+    input_dataset = "none-yet/anime-captions"
+    output_dataset = "none-yet/anime-captions"
+    model_name = "datalab-to/chandra"
 
-loaded = datasets.load_dataset(input_dataset, split="train")
+    loaded = datasets.load_dataset(input_dataset, split="train")
 
-if isinstance(loaded, datasets.DatasetDict):
-    ds = cast(Dataset, loaded["train"])
-else:
-    ds = cast(Dataset, loaded)
+    if isinstance(loaded, datasets.DatasetDict):
+        ds = cast(Dataset, loaded["train"])
+    else:
+        ds = cast(Dataset, loaded)
 
-num_gpus = torch.cuda.device_count()
-batch_size = 32
-total_size = len(ds)
-shard_size = total_size // num_gpus
+    num_gpus = torch.cuda.device_count()
+    batch_size = 8
+    total_size = len(ds)
+    shard_size = total_size // num_gpus
 
-worker_args = []
-for i in range(num_gpus):
-    start = i * shard_size
-    end = start + shard_size if i < num_gpus - 1 else total_size
-    worker_args.append((i, i, start, end, model_name, batch_size))
+    worker_args = []
+    for i in range(num_gpus):
+        start = i * shard_size
+        end = start + shard_size if i < num_gpus - 1 else total_size
+        worker_args.append((i, i, start, end, model_name, batch_size))
 
-with mp.Pool(processes=num_gpus) as pool:
-    shards = pool.map(process_shard_worker, worker_args)
+    with mp.Pool(processes=num_gpus) as pool:
+        shards = pool.map(process_shard_worker, worker_args)
 
-ds = datasets.concatenate_datasets(shards)
+    ds = datasets.concatenate_datasets(shards)
 
 # %%
 ds.push_to_hub(output_dataset)