def main(opts): """ Opts prevalence: 1. Load file specified in args.default (or shared/trainer/defaults.yaml if none is provided) 2. Update with file specified in args.config (or no update if none is provided) 3. Update with parsed command-line arguments e.g. `python train.py args.config=config/large-lr.yaml data.loaders.batch_size=10` loads defaults, overrides with values in large-lr.yaml and sets batch_size to 10 """ # ----------------------------- # ----- Parse arguments ----- # ----------------------------- hydra_opts = Dict(OmegaConf.to_container(opts)) args = hydra_opts.pop("args", None) auto_resumed = {} config_path = args.config if hydra_opts.train.resume: out_ = str(env_to_path(hydra_opts.output_path)) config_path = Path(out_) / "opts.yaml" if not config_path.exists(): config_path = None print("WARNING: could not reuse the opts in {}".format(out_)) default = args.default or Path( __file__).parent / "shared/trainer/defaults.yaml" # ----------------------- # ----- Load opts ----- # ----------------------- opts = load_opts(config_path, default=default, commandline_opts=hydra_opts) if args.resume: opts.train.resume = True opts.jobID = os.environ.get("SLURM_JOBID") opts.slurm_partition = os.environ.get("SLURM_JOB_PARTITION") opts.output_path = str(env_to_path(opts.output_path)) print("Config output_path:", opts.output_path) exp = comet_previous_id = None # ------------------------------- # ----- Check output_path ----- # ------------------------------- # Auto-continue if same slurm job ID (=job was requeued) if not opts.train.resume and opts.train.auto_resume: print("\n\nTrying to auto-resume...") existing_path = find_existing_training(opts) if existing_path is not None and existing_path.exists(): auto_resumed["original output_path"] = str(opts.output_path) auto_resumed["existing_path"] = str(existing_path) opts.train.resume = True opts.output_path = str(existing_path) # Still not resuming: creating new output path if not opts.train.resume: opts.output_path = str(get_increased_path(opts.output_path)) Path(opts.output_path).mkdir(parents=True, exist_ok=True) # Copy the opts's sbatch_file to output_path copy_run_files(opts) # store git hash opts.git_hash = get_git_revision_hash() opts.git_branch = get_git_branch() if not args.no_comet: # ---------------------------------- # ----- Set Comet Experiment ----- # ---------------------------------- if opts.train.resume: # Is resuming: get existing comet exp id assert Path( opts.output_path).exists(), "Output_path does not exist" comet_previous_id = get_existing_comet_id(opts.output_path) # Continue existing experiment if comet_previous_id is None: print("WARNING could not retreive previous comet id") print(f"from {opts.output_path}") else: print("Continuing previous experiment", comet_previous_id) auto_resumed["continuing exp id"] = comet_previous_id exp = ExistingExperiment(previous_experiment=comet_previous_id, **comet_kwargs) print("Comet Experiment resumed") if exp is None: # Create new experiment print("Starting new experiment") exp = Experiment(project_name="climategan", **comet_kwargs) exp.log_asset_folder( str(Path(__file__).parent / "climategan"), recursive=True, log_file_name=True, ) exp.log_asset(str(Path(__file__))) # Log note if args.note: exp.log_parameter("note", args.note) # Merge and log tags if args.comet_tags or opts.comet.tags: tags = set([f"branch:{opts.git_branch}"]) if args.comet_tags: tags.update(args.comet_tags) if opts.comet.tags: tags.update(opts.comet.tags) opts.comet.tags = list(tags) print("Logging to comet.ml with tags", opts.comet.tags) exp.add_tags(opts.comet.tags) # Log all opts exp.log_parameters(flatten_opts(opts)) if auto_resumed: exp.log_text("\n".join(f"{k:20}: {v}" for k, v in auto_resumed.items())) # allow some time for comet to get its url sleep(1) # Save comet exp url url_path = get_increased_path(Path(opts.output_path) / "comet_url.txt") with open(url_path, "w") as f: f.write(exp.url) # Save config file opts_path = get_increased_path(Path(opts.output_path) / "opts.yaml") with (opts_path).open("w") as f: yaml.safe_dump(opts.to_dict(), f) pprint("Running model in", opts.output_path) # ------------------- # ----- Train ----- # ------------------- trainer = Trainer(opts, comet_exp=exp, verbose=1) trainer.logger.time.start_time = time() trainer.setup() trainer.train() # ----------------------------- # ----- End of training ----- # ----------------------------- pprint("Done training") kill_job(opts.jobID)
class CometConnection: def __init__(self, comet_name=None, dataset_config=None, exp_key=None): self.experiment = None if comet_name is not None and dataset_config is not None: self._init_new_experiment(comet_name, dataset_config) elif exp_key is not None: self._init_continue_experiment(exp_key) def _init_new_experiment(self, comet_name, dataset_config): self.experiment = Experiment(api_key=COMET_KEY, project_name=PROJECT_NAME) self.experiment.set_name(comet_name) self.log_data_attributes(dataset_config) self.experiment.log_asset('datagen/spectra_generator.m') def _init_continue_experiment(self, exp_key): self.experiment = ExistingExperiment(api_key=COMET_KEY, previous_experiment=exp_key) def serialize(self): params = dict() params["comet_exp_key"] = self.experiment.get_key() return params def save(self, save_dir): info_dict = self.serialize() json.dump(info_dict, open(os.path.join(save_dir, COMET_SAVE_FILENAME), "w")) def persist(self, config_path): info = json.load(open(config_path, 'r')) self.__init__(exp_key=info["comet_exp_key"]) def log_data_attributes(self, dataset_config): for key, value in dataset_config.items(): self.experiment.log_parameter("SPECTRUM_" + key, value) def log_imgs(self, dataset_name): try: imgs_dir = os.path.join(DATA_DIR, dataset_name, 'imgs') self.experiment.log_asset_folder(imgs_dir) except: print(f"No images found for dataset: {dataset_name}") def log_script(self, dataset_config): script_name = dataset_config['matlab_script'] try: matlab_dir = os.path.join(GEN_DIR, script_name) self.experiment.log_asset(matlab_dir) except: print(f"Could not find {script_name} under {GEN_DIR}.") def format_classification_report(self, classification_report): return { f'{k}_test_{metric}': metric_val for k, v in classification_report.items() for metric, metric_val in v.items() } def get_classification_report(self, y_test, preds): preds_formatted = np.argmax(preds, axis=1) test_formatted = np.argmax(y_test, axis=1) peak_labels = [ f"n_peaks_{1 + num_peak}" for num_peak in range(y_test.shape[1]) ] classif_report = classification_report(test_formatted, preds_formatted, target_names=peak_labels, output_dict=True) classif_report_str = classification_report(test_formatted, preds_formatted, target_names=peak_labels) if self.experiment is not None: formatted = self.format_classification_report(classif_report) self.experiment.log_metrics(formatted) self.experiment.log_text(classif_report_str) return classif_report