| |
| |
|
|
| |
| |
|
|
| import os |
| from train import main |
| from submitit.helpers import Checkpointable |
|
|
| LOCAL = False |
| try: |
| import submitit |
| except: |
| print( |
| "No submitit package found! Defaulting to executing the script in the local machine" |
| ) |
| LOCAL = True |
| import parser |
| import json |
|
|
|
|
| class Trainer(Checkpointable): |
| def __call__(self, args, slurm=False): |
| if slurm and not LOCAL: |
| main( |
| args, |
| args.outdir, |
| master_node=submitit.JobEnvironment().hostnames[0], |
| port=args.port, |
| ) |
| else: |
| main(args, args.outdir, master_node="", dry_run=args.dry_run) |
|
|
|
|
| if __name__ == "__main__": |
| parser_ = parser.get_parser() |
| args = parser_.parse_args() |
|
|
| if args.json_config != "": |
| data = json.load(open(args.json_config)) |
| for key in data.keys(): |
| setattr(args, key, data[key]) |
| else: |
| print("Not using JSON configuration file!") |
| if args.data_root is not None: |
| print("Appending data_root to paths") |
| args.data = os.path.join(args.data_root, args.data) |
| args.root_feats = os.path.join(args.data_root, args.root_feats) |
| args.root_nns = os.path.join(args.data_root, args.root_nns) |
| args.outdir = args.base_root |
|
|
| trainer = Trainer() |
| if not args.slurm or LOCAL: |
| trainer(args) |
| else: |
|
|
| executor = submitit.SlurmExecutor(folder=args.slurm_logdir, max_num_timeout=60) |
| print(args.gpus) |
| executor.update_parameters( |
| gpus_per_node=args.gpus, |
| partition=args.partition, |
| constraint="volta32gb", |
| nodes=args.nodes, |
| ntasks_per_node=args.gpus, |
| cpus_per_task=10, |
| mem=256000, |
| time=args.slurm_time, |
| job_name=args.exp_name, |
| exclusive=True if args.gpus == 8 else False, |
| ) |
|
|
| job = executor.submit(trainer, args, slurm=True) |
| print(job.job_id) |
|
|
| import time |
|
|
| time.sleep(1) |
|
|
