Spaces:
Running
Running
File size: 1,714 Bytes
c4c7cee |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 |
import sys
from pathlib import Path
sys.path.append(str(Path(__file__).resolve().parent.parent.parent))
import argparse
import os
from jean_zay.launch import JeanZayExperiment
def parse_mode():
parser = argparse.ArgumentParser(
description="Extract embeddings from a dataset using DINOv2"
)
parser.add_argument(
"--launch",
action="store_true",
help="Launch the experiment",
)
parser.add_argument(
"--number_of_splits",
type=int,
help="Number of splits to process",
default=1,
)
parser.add_argument(
"--input_path",
type=str,
help="Path to the input dataset",
)
parser.add_argument(
"--output_path",
type=str,
help="Path to the output dataset",
)
args = parser.parse_args()
return args
args = parse_mode()
cmd_modifiers = []
exps = []
exp_name = f"preprocess_data"
job_name = f"preprocess_data"
jz_exp = JeanZayExperiment(
exp_name,
job_name,
slurm_array_nb_jobs=args.number_of_splits,
cmd_path="data/extract_embeddings/dino_v2.py",
num_nodes=1,
num_gpus_per_node=1,
qos="t3",
account="mya",
gpu_type="h100",
time="02:00:00",
)
exps.append(jz_exp)
trainer_modifiers = {}
exp_modifier = {
"--input_path": args.input_path,
"--output_path": args.output_path,
"--number_of_splits": args.number_of_splits,
"--split_index": "${SLURM_ARRAY_TASK_ID}",
}
cmd_modifiers.append(dict(trainer_modifiers, **exp_modifier))
if __name__ == "__main__":
for exp, cmd_modifier in zip(exps, cmd_modifiers):
exp.build_cmd(cmd_modifier)
if args.launch == True:
exp.launch()
|