class CometLogger(): def __init__(self, enabled, is_existing=False, prev_exp_key=None): """ Handles logging of experiment to comet and also persistence to local file system. Supports resumption of stopped experiments. """ disabled = not enabled if not is_existing: self.experiment = Experiment(api_key=COMET_API_KEY, workspace=COMET_WORKSPACE, project_name=PROJECT_NAME, disabled=disabled) else: if prev_exp_key is None: raise ValueError( "Requested existing experiment, but no key provided") print("Continuing existing experiment with key: ", prev_exp_key) self.experiment = ExistingExperiment( api_key=COMET_API_KEY, workspace=COMET_WORKSPACE, project_name=PROJECT_NAME, disabled=disabled, previous_experiment=prev_exp_key) self.disabled = disabled def get_experiment_key(self): return self.experiment.get_key()[:9] def add_tag(self, tag): self.experiment.add_tag(tag) def log_metric(self, name, value, step=None): self.experiment.log_metric(name, value, step=step) def log_metrics(self, metrics_dict, prefix, step=None): self.experiment.log_metrics(metrics_dict, prefix=prefix, step=step) def log_params(self, params_dict): self.experiment.log_parameters(params_dict) def set_name(self, name_str): self.experiment.set_name(name_str) def log_dataset(self, dataset: SpeakerVerificationDataset): if self.disabled: return dataset_string = "" dataset_string += "<b>Speakers</b>: %s\n" % len(dataset.speakers) dataset_string += "\n" + dataset.get_logs() dataset_string = dataset_string.replace("\n", "<br>") self.vis.text(dataset_string, opts={"title": "Dataset"}) def log_implementation(self, params): if self.disabled: return implementation_string = "" for param, value in params.items(): implementation_string += "<b>%s</b>: %s\n" % (param, value) implementation_string = implementation_string.replace("\n", "<br>") self.implementation_string = implementation_string self.implementation_win = self.vis.text( implementation_string, opts={"title": "Training implementation"}) def draw_projections(self, embeds, utterances_per_speaker, step, out_fpath=None, max_speakers=16): if self.disabled: return max_speakers = min(max_speakers, len(colormap)) embeds = embeds[:max_speakers * utterances_per_speaker] n_speakers = len(embeds) // utterances_per_speaker ground_truth = np.repeat(np.arange(n_speakers), utterances_per_speaker) colors = [colormap[i] for i in ground_truth] reducer = umap.UMAP() projected = reducer.fit_transform(embeds) plt.scatter(projected[:, 0], projected[:, 1], c=colors) plt.gca().set_aspect("equal", "datalim") plt.title("UMAP projection (step %d)" % step) if out_fpath is not None: plt.savefig(out_fpath) plt.clf() self.experiment.log_image(out_fpath, step=step)
def main(args, config=None, init_distributed=False): utils.import_user_module(args) experiment = None if config: experiment = ExistingExperiment( api_key=config["api_key"], previous_experiment=config["experiment_key"], auto_output_logging=None, ) assert ( args.max_tokens is not None or args.max_sentences is not None ), "Must specify batch size either with --max-tokens or --max-sentences" # Initialize CUDA and distributed training if torch.cuda.is_available() and not args.cpu: torch.cuda.set_device(args.device_id) np.random.seed(args.seed) torch.manual_seed(args.seed) if init_distributed: args.distributed_rank = distributed_utils.distributed_init(args) if distributed_utils.is_master(args): checkpoint_utils.verify_checkpoint_directory(args.save_dir) print(args) if experiment: experiment.log_parameters(vars(args), prefix="Device {} :: ".format( args.device_id)) # Setup task, e.g., translation, language modeling, etc. task = tasks.setup_task(args) # Load valid dataset (we load training data below, based on the latest checkpoint) for valid_sub_split in args.valid_subset.split(","): task.load_dataset(valid_sub_split, combine=False, epoch=0) # Build model and criterion model = task.build_model(args) criterion = task.build_criterion(args) print(model) print("| model {}, criterion {}".format(args.arch, criterion.__class__.__name__)) print("| num. model params: {} (num. trained: {})".format( sum(p.numel() for p in model.parameters()), sum(p.numel() for p in model.parameters() if p.requires_grad), )) if experiment: experiment.log_parameters( { "criterion": criterion.__class__.__name__, "num. model params": sum(p.numel() for p in model.parameters()), "num. trained params": sum(p.numel() for p in model.parameters() if p.requires_grad), }, prefix="Device {} :: ".format(args.device_id), ) # Build trainer trainer = Trainer(args, task, model, criterion) print("| training on {} GPUs".format(args.distributed_world_size)) print("| max tokens per GPU = {} and max sentences per GPU = {}".format( args.max_tokens, args.max_sentences)) # Load the latest checkpoint if one is available and restore the # corresponding train iterator extra_state, epoch_itr = checkpoint_utils.load_checkpoint(args, trainer) # Train until the learning rate gets too small max_epoch = args.max_epoch or math.inf max_update = args.max_update or math.inf lr = trainer.get_lr() train_meter = StopwatchMeter() train_meter.start() valid_subsets = args.valid_subset.split(",") while (lr > args.min_lr and epoch_itr.epoch < max_epoch and trainer.get_num_updates() < max_update): # train for one epoch train(args, trainer, task, epoch_itr, experiment) if (not args.disable_validation and epoch_itr.epoch % args.validate_interval == 0): valid_losses = validate(args, trainer, task, epoch_itr, valid_subsets, experiment) else: valid_losses = [None] # only use first validation loss to update the learning rate lr = trainer.lr_step(epoch_itr.epoch, valid_losses[0]) # save checkpoint if epoch_itr.epoch % args.save_interval == 0: checkpoint_utils.save_checkpoint(args, trainer, epoch_itr, valid_losses[0]) reload_dataset = ":" in getattr(args, "data", "") # sharded data: get train iterator for next epoch epoch_itr = trainer.get_train_iterator(epoch_itr.epoch, load_dataset=reload_dataset) train_meter.stop() print("| done training in {:.1f} seconds".format(train_meter.sum)) if experiment: experiment.log_metrics( { "valid_loss": valid_losses[0], "lr": lr }, prefix="Device {} ".format(args.device_id), )
class CometLogger(): def __init__(self, disabled, is_existing=False, prev_exp_key=None): """ Handles logging of experiment to comet and also persistence to local file system. Supports resumption of stopped experiments. """ if not is_existing: self.experiment = Experiment(api_key=COMET_API_KEY, workspace=COMET_WORKSPACE, project_name=PROJECT_NAME, disabled=disabled) else: if prev_exp_key is None: raise ValueError("Requested existing experiment, but no key provided") print("Continuing existing experiment with key: ", prev_exp_key) self.experiment = ExistingExperiment(api_key=COMET_API_KEY, workspace=COMET_WORKSPACE, project_name=PROJECT_NAME, disabled=disabled, previous_experiment=prev_exp_key) self.disabled = disabled self.name = None def get_experiment_key(self): return self.experiment.get_key()[:9] def add_tag(self, tag): self.experiment.add_tag(tag) def log_metric(self, name, value, step=None): self.experiment.log_metric(name, value, step=step) def log_metrics(self, metrics_dict, prefix, step=None): self.experiment.log_metrics(metrics_dict, prefix=prefix, step=step) def log_params(self, params_dict): self.experiment.log_parameters(params_dict) def set_name(self, name_str): self.experiment.set_name(name_str) self.name = name_str def save_act_grads(self, log_dict): """Save a dictionary of activation/gradients records to disk""" assert isinstance(log_dict, dict) if self.name is None: warnings.warn("Experiment name not set, not saving") return # Save the log dictionary file_name = f"./.{self.name}.record" with open(file_name, 'wb') as f: pickle.dump(log_dict, f) # TODO: need to rewrite before can be used for MNIST. def draw_projections(self, embeds, utterances_per_speaker, step, out_fpath=None, max_speakers=16): import umap if self.disabled: return max_speakers = min(max_speakers, len(colormap)) embeds = embeds[:max_speakers * utterances_per_speaker] n_speakers = len(embeds) // utterances_per_speaker ground_truth = np.repeat(np.arange(n_speakers), utterances_per_speaker) colors = [colormap[i] for i in ground_truth] reducer = umap.UMAP() projected = reducer.fit_transform(embeds) plt.scatter(projected[:, 0], projected[:, 1], c=colors) plt.gca().set_aspect("equal", "datalim") plt.title("UMAP projection (step %d)" % step) if out_fpath is not None: plt.savefig(out_fpath) plt.clf() self.experiment.log_image(out_fpath, step=step)
valSSIM[epoch].append(ssim_val) valLoss[epoch].append(vLoss) # Tensorboard itr = int(trainIndex + epoch * (len(trainloader))) # writer.add_scalars( # "Loss", # {"trainLoss": iLoss / args.progress_iter, "validationLoss": vLoss}, # itr, # ) # writer.add_scalar("PSNR", psnr, itr) # writer.add_image("Validation", valImg, itr) comet_exp.log_metrics( {"trainLoss": iLoss / args.progress_iter, "validationLoss": vLoss}, step=itr, epoch=epoch, ) comet_exp.log_metric("PSNR", psnr, step=itr, epoch=epoch) comet_exp.log_metric("SSIM", ssim_val, step=itr, epoch=epoch) # valImage = torch.movedim(valImg, 0, -1) # print(type(valImage)) # print(valImage.shape) # print(valImage.max()) # print(valImage.min()) # comet_exp.log_image( # valImage, # name="iter: " + str(iter) + ";epoch: " + str(epoch), # image_format="jpg", # step=itr,
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