class TensorboardLogger: def __init__(self, log_folder="./logs", iteration=0): # This would handle warning of missing tensorboard from torch.utils.tensorboard import SummaryWriter self.summary_writer = None self._is_master = is_master() self.timer = Timer() self.log_folder = log_folder self.time_format = "%Y-%m-%dT%H:%M:%S" if self._is_master: current_time = self.timer.get_time_hhmmss(None, format=self.time_format) tensorboard_folder = os.path.join( self.log_folder, f"tensorboard_{current_time}" ) self.summary_writer = SummaryWriter(tensorboard_folder) def __del__(self): if getattr(self, "summary_writer", None) is not None: self.summary_writer.close() def _should_log_tensorboard(self): if self.summary_writer is None or not self._is_master: return False else: return True def add_scalar(self, key, value, iteration): if not self._should_log_tensorboard(): return self.summary_writer.add_scalar(key, value, iteration) def add_scalars(self, scalar_dict, iteration): if not self._should_log_tensorboard(): return for key, val in scalar_dict.items(): self.summary_writer.add_scalar(key, val, iteration) def add_histogram_for_model(self, model, iteration): if not self._should_log_tensorboard(): return for name, param in model.named_parameters(): np_param = param.clone().cpu().data.numpy() self.summary_writer.add_histogram(name, np_param, iteration) def flush(self): if self._should_log_tensorboard(): self.summary_writer.flush()
class TensorboardLogger: def __init__(self, log_folder="./logs", iteration=0): self._summary_writer = None self._is_main = is_main() self.timer = Timer() self.log_folder = log_folder self.time_format = "%Y-%m-%dT%H:%M:%S" current_time = self.timer.get_time_hhmmss(None, format=self.time_format) self.tensorboard_folder = os.path.join(self.log_folder, f"tensorboard_{current_time}") @property def summary_writer(self): # Only on rank zero if not self._is_main: return None if self._summary_writer is None: # This would handle warning of missing tensorboard from torch.utils.tensorboard import SummaryWriter self._summary_writer = SummaryWriter(self.tensorboard_folder) return self._summary_writer @skip_if_tensorboard_inactive def close(self): """ Closes the tensorboard summary writer. """ self.summary_writer.close() @skip_if_tensorboard_inactive def add_scalar(self, key, value, iteration): self.summary_writer.add_scalar(key, value, iteration) @skip_if_tensorboard_inactive def add_scalars(self, scalar_dict, iteration): for key, val in scalar_dict.items(): self.summary_writer.add_scalar(key, val, iteration) @skip_if_tensorboard_inactive def add_histogram_for_model(self, model, iteration): for name, param in model.named_parameters(): np_param = param.clone().cpu().data.numpy() self.summary_writer.add_histogram(name, np_param, iteration)
class LogisticsCallback(Callback): """Callback for handling train/validation logistics, report summarization, logging etc. """ def __init__(self, config, trainer): """ Attr: config(mmf_typings.DictConfig): Config for the callback trainer(Type[BaseTrainer]): Trainer object """ super().__init__(config, trainer) self.total_timer = Timer() self.log_interval = self.training_config.log_interval self.evaluation_interval = self.training_config.evaluation_interval self.checkpoint_interval = self.training_config.checkpoint_interval # Total iterations for snapshot self.snapshot_iterations = len(self.trainer.val_dataset) self.snapshot_iterations //= self.training_config.batch_size self.tb_writer = None if self.training_config.tensorboard: log_dir = setup_output_folder(folder_only=True) env_tb_logdir = get_mmf_env(key="tensorboard_logdir") if env_tb_logdir: log_dir = env_tb_logdir self.tb_writer = TensorboardLogger(log_dir, self.trainer.current_iteration) def on_train_start(self): self.train_timer = Timer() self.snapshot_timer = Timer() def on_update_end(self, **kwargs): if not kwargs["should_log"]: return extra = {} if "cuda" in str(self.trainer.device): extra["max mem"] = torch.cuda.max_memory_allocated() / 1024 extra["max mem"] //= 1024 if self.training_config.experiment_name: extra["experiment"] = self.training_config.experiment_name extra.update( { "epoch": self.trainer.current_epoch, "num_updates": self.trainer.num_updates, "iterations": self.trainer.current_iteration, "max_updates": self.trainer.max_updates, "lr": "{:.5f}".format( self.trainer.optimizer.param_groups[0]["lr"] ).rstrip("0"), "ups": "{:.2f}".format( self.log_interval / self.train_timer.unix_time_since_start() ), "time": self.train_timer.get_time_since_start(), "time_since_start": self.total_timer.get_time_since_start(), "eta": self._calculate_time_left(), } ) self.train_timer.reset() self._summarize_report(kwargs["meter"], extra=extra) def on_validation_start(self, **kwargs): self.snapshot_timer.reset() def on_validation_end(self, **kwargs): extra = { "num_updates": self.trainer.num_updates, "epoch": self.trainer.current_epoch, "iterations": self.trainer.current_iteration, "max_updates": self.trainer.max_updates, "val_time": self.snapshot_timer.get_time_since_start(), } extra.update(self.trainer.early_stop_callback.early_stopping.get_info()) self.train_timer.reset() self._summarize_report(kwargs["meter"], extra=extra) def on_test_end(self, **kwargs): prefix = "{}: full {}".format( kwargs["report"].dataset_name, kwargs["report"].dataset_type ) self._summarize_report(kwargs["meter"], prefix) logger.info(f"Finished run in {self.total_timer.get_time_since_start()}") def _summarize_report(self, meter, should_print=True, extra=None): if extra is None: extra = {} if not is_master() and not is_xla(): return if self.training_config.tensorboard: scalar_dict = meter.get_scalar_dict() self.tb_writer.add_scalars(scalar_dict, self.trainer.current_iteration) if not should_print: return log_dict = {} if hasattr(self.trainer, "num_updates") and hasattr( self.trainer, "max_updates" ): log_dict.update( {"progress": f"{self.trainer.num_updates}/{self.trainer.max_updates}"} ) log_dict.update(meter.get_log_dict()) log_dict.update(extra) log_progress(log_dict) def _calculate_time_left(self): time_taken_for_log = time.time() * 1000 - self.train_timer.start iterations_left = self.trainer.max_updates - self.trainer.num_updates num_logs_left = iterations_left / self.log_interval time_left = num_logs_left * time_taken_for_log snapshot_iteration = self.snapshot_iterations / self.log_interval snapshot_iteration *= iterations_left / self.evaluation_interval time_left += snapshot_iteration * time_taken_for_log return self.train_timer.get_time_hhmmss(gap=time_left)
class BaseTrainer: def __init__(self, configuration): self.configuration = configuration self.config = self.configuration.get_config() self.profiler = Timer() self.total_timer = Timer() if self.configuration is not None: self.args = self.configuration.args def load(self): self._set_device() self.run_type = self.config.get("run_type", "train") self.dataset_loader = DatasetLoader(self.config) self._datasets = self.config.datasets # Check if loader is already defined, else init it writer = registry.get("writer", no_warning=True) if writer: self.writer = writer else: self.writer = Logger(self.config) registry.register("writer", self.writer) self.configuration.pretty_print() self.config_based_setup() self.load_datasets() self.load_model_and_optimizer() self.load_metrics() def _set_device(self): self.local_rank = self.config.device_id self.device = self.local_rank self.distributed = False # Will be updated later based on distributed setup registry.register("global_device", self.device) if self.config.distributed.init_method is not None: self.distributed = True self.device = torch.device("cuda", self.local_rank) elif torch.cuda.is_available(): self.device = torch.device("cuda") else: self.device = torch.device("cpu") registry.register("current_device", self.device) def load_datasets(self): self.writer.write("Loading datasets", "info") self.dataset_loader.load_datasets() self.train_dataset = self.dataset_loader.train_dataset self.val_dataset = self.dataset_loader.val_dataset # Total iterations for snapshot self.snapshot_iterations = len(self.val_dataset) self.snapshot_iterations //= self.config.training.batch_size self.test_dataset = self.dataset_loader.test_dataset self.train_loader = self.dataset_loader.train_loader self.val_loader = self.dataset_loader.val_loader self.test_loader = self.dataset_loader.test_loader def load_metrics(self): metrics = self.config.evaluation.get("metrics", []) self.metrics = Metrics(metrics) self.metrics_params = self.metrics.required_params def load_model_and_optimizer(self): attributes = self.config.model_config[self.config.model] # Easy way to point to config for other model if isinstance(attributes, str): attributes = self.config.model_config[attributes] with omegaconf.open_dict(attributes): attributes.model = self.config.model self.model = build_model(attributes) if "cuda" in str(self.device): device_info = "CUDA Device {} is: {}".format( self.config.distributed.rank, torch.cuda.get_device_name(self.local_rank), ) registry.register("global_device", self.config.distributed.rank) self.writer.write(device_info, log_all=True) self.model = self.model.to(self.device) self.optimizer = build_optimizer(self.model, self.config) registry.register("data_parallel", False) registry.register("distributed", False) self.load_extras() self.parallelize_model() def parallelize_model(self): training = self.config.training if ("cuda" in str(self.device) and torch.cuda.device_count() > 1 and not self.distributed): registry.register("data_parallel", True) self.model = torch.nn.DataParallel(self.model) if "cuda" in str(self.device) and self.distributed: registry.register("distributed", True) self.model = torch.nn.parallel.DistributedDataParallel( self.model, device_ids=[self.local_rank], output_device=self.local_rank, check_reduction=True, find_unused_parameters=training.find_unused_parameters, ) def load_extras(self): self.writer.write("Torch version is: " + torch.__version__) self.checkpoint = Checkpoint(self) self.meter = Meter() self.training_config = self.config.training early_stop_criteria = self.training_config.early_stop.criteria early_stop_minimize = self.training_config.early_stop.minimize early_stop_enabled = self.training_config.early_stop.enabled early_stop_patience = self.training_config.early_stop.patience self.log_interval = self.training_config.log_interval self.evaluation_interval = self.training_config.evaluation_interval self.checkpoint_interval = self.training_config.checkpoint_interval self.max_updates = self.training_config.max_updates self.should_clip_gradients = self.training_config.clip_gradients self.max_epochs = self.training_config.max_epochs self.early_stopping = EarlyStopping( self.model, self.checkpoint, early_stop_criteria, patience=early_stop_patience, minimize=early_stop_minimize, should_stop=early_stop_enabled, ) self.current_epoch = 0 self.current_iteration = 0 self.num_updates = 0 self.checkpoint.load_state_dict() self.not_debug = self.training_config.logger_level != "debug" self.lr_scheduler = None if self.training_config.lr_scheduler is True: self.lr_scheduler = build_scheduler(self.optimizer, self.config) self.tb_writer = None if self.training_config.tensorboard: log_dir = self.writer.log_dir env_tb_logdir = get_mmf_env(key="tensorboard_logdir") if env_tb_logdir: log_dir = env_tb_logdir self.tb_writer = TensorboardLogger(log_dir, self.current_iteration) def config_based_setup(self): seed = self.config.training.seed if seed is None: return torch.backends.cudnn.deterministic = True torch.backends.cudnn.benchmark = False def train(self): self.writer.write("===== Model =====") self.writer.write(self.model) print_model_parameters(self.model) if "train" not in self.run_type: self.inference() return should_break = False if self.max_epochs is None: self.max_epochs = math.inf else: self.max_updates = math.inf self.model.train() self.train_timer = Timer() self.snapshot_timer = Timer() self.profile("Setup Time") torch.autograd.set_detect_anomaly(True) self.writer.write("Starting training...") while self.num_updates < self.max_updates and not should_break: self.current_epoch += 1 registry.register("current_epoch", self.current_epoch) # Seed the sampler in case if it is distributed self.dataset_loader.seed_sampler("train", self.current_epoch) if self.current_epoch > self.max_epochs: break for batch in self.train_loader: self.profile("Batch load time") self.current_iteration += 1 self.writer.write(self.num_updates + 1, "debug") report = self._forward_pass(batch) loss = self._extract_loss(report) self._backward(loss) should_break = self._logistics(report) if self.num_updates > self.max_updates: should_break = True if should_break: break # In distributed, each worker will complete one epoch when we reach this # as each worker is an individual instance self.current_epoch += get_world_size() - 1 self.finalize() def _run_scheduler(self): if self.lr_scheduler is not None: self.lr_scheduler.step(self.num_updates) def _forward_pass(self, batch): prepared_batch = self.dataset_loader.prepare_batch(batch) self.profile("Batch prepare time") # Arguments should be a dict at this point model_output = self.model(prepared_batch) report = Report(prepared_batch, model_output) self.profile("Forward time") return report def _backward(self, loss): self.optimizer.zero_grad() loss.backward() if self.should_clip_gradients: clip_gradients(self.model, self.num_updates, self.tb_writer, self.config) self.optimizer.step() self._run_scheduler() self.num_updates += 1 self.profile("Backward time") def _extract_loss(self, report): loss_dict = report.losses loss = sum([loss.mean() for loss in loss_dict.values()]) return loss def finalize(self): self.writer.write("Stepping into final validation check") # Only do when run_type has train as it shouldn't happen on validation and # inference runs. Inference will take care of this anyways. Also, don't run # if current iteration is divisble by snapshot interval as it will just # be a repeat if ("train" in self.run_type and self.num_updates % self.evaluation_interval != 0): self._try_full_validation(force=True) self.checkpoint.restore() self.checkpoint.finalize() self.inference() self.writer.write( f"Finished run in {self.total_timer.get_time_since_start()}") def _update_meter(self, report, meter=None, eval_mode=False): if meter is None: meter = self.meter if hasattr(report, "metrics"): metrics_dict = report.metrics reduced_metrics_dict = reduce_dict(metrics_dict) if not eval_mode: loss_dict = report.losses reduced_loss_dict = reduce_dict(loss_dict) with torch.no_grad(): # Add metrics to meter only when mode is `eval` meter_update_dict = {} if not eval_mode: loss_key = report.dataset_type + "/total_loss" reduced_loss = sum( [loss.mean() for loss in reduced_loss_dict.values()]) if hasattr(reduced_loss, "item"): reduced_loss = reduced_loss.item() registry.register(loss_key, reduced_loss) meter_update_dict.update({loss_key: reduced_loss}) meter_update_dict.update(reduced_loss_dict) if hasattr(report, "metrics"): meter_update_dict.update(reduced_metrics_dict) meter.update(meter_update_dict, report.batch_size) def _logistics(self, report): registry.register("current_iteration", self.current_iteration) registry.register("num_updates", self.num_updates) should_print = self.num_updates % self.log_interval == 0 should_break = False extra = {} if should_print is True: if "cuda" in str(self.device): extra["max mem"] = torch.cuda.max_memory_allocated() / 1024 extra["max mem"] //= 1024 if self.training_config.experiment_name: extra["experiment"] = self.training_config.experiment_name extra.update({ "epoch": self.current_epoch, "num_updates": self.num_updates, "iterations": self.current_iteration, "max_updates": self.max_updates, "lr": "{:.5f}".format( self.optimizer.param_groups[0]["lr"]).rstrip("0"), "ups": "{:.2f}".format(self.log_interval / self.train_timer.unix_time_since_start()), "time": self.train_timer.get_time_since_start(), "time_since_start": self.total_timer.get_time_since_start(), "eta": self._calculate_time_left(), }) self.train_timer.reset() # Calculate metrics every log interval for debugging if self.training_config.evaluate_metrics: report.metrics = self.metrics(report, report) self._update_meter(report, self.meter) self._summarize_report(self.meter, should_print=should_print, extra=extra) self._try_snapshot() should_break = self._try_full_validation() return should_break def _try_snapshot(self): if self.num_updates % self.checkpoint_interval == 0: self.writer.write("Checkpoint time. Saving a checkpoint.") self.checkpoint.save(self.num_updates, self.current_iteration, update_best=False) def _try_full_validation(self, force=False): should_break = False if self.num_updates % self.evaluation_interval == 0 or force: self.snapshot_timer.reset() self.writer.write( "Evaluation time. Running on full validation set...") # Validation and Early stopping # Create a new meter for this case report, meter = self.evaluate(self.val_loader) extra = { "num_updates": self.num_updates, "epoch": self.current_epoch, "iterations": self.current_iteration, "max_updates": self.max_updates, "val_time": self.snapshot_timer.get_time_since_start(), } stop = self.early_stopping(self.num_updates, self.current_iteration, meter) stop = bool(broadcast_scalar(stop, src=0, device=self.device)) extra.update(self.early_stopping.get_info()) self._summarize_report(meter, extra=extra) gc.collect() if "cuda" in str(self.device): torch.cuda.empty_cache() if stop is True: self.writer.write("Early stopping activated") should_break = True self.train_timer.reset() return should_break def evaluate(self, loader, use_tqdm=False, single_batch=False): meter = Meter() with torch.no_grad(): self.model.eval() disable_tqdm = not use_tqdm or not is_master() combined_report = None for batch in tqdm(loader, disable=disable_tqdm): report = self._forward_pass(batch) self._update_meter(report, meter) # accumulate necessary params for metric calculation if combined_report is None: combined_report = report else: combined_report.accumulate_tensor_fields( report, self.metrics.required_params) combined_report.batch_size += report.batch_size if single_batch is True: break combined_report.metrics = self.metrics(combined_report, combined_report) self._update_meter(combined_report, meter, eval_mode=True) self.model.train() return combined_report, meter def _summarize_report(self, meter, should_print=True, extra=None): if extra is None: extra = {} if not is_master(): return if self.training_config.tensorboard: scalar_dict = meter.get_scalar_dict() self.tb_writer.add_scalars(scalar_dict, self.current_iteration) if not should_print: return log_dict = {"progress": f"{self.num_updates}/{self.max_updates}"} log_dict.update(meter.get_log_dict()) log_dict.update(extra) self.writer.log_progress(log_dict) def inference(self): if "val" in self.run_type: self._inference_run("val") if any(rt in self.run_type for rt in ["inference", "test", "predict"]): self._inference_run("test") def _inference_run(self, dataset_type): if self.config.evaluation.predict: self.predict(dataset_type) return self.writer.write(f"Starting inference on {dataset_type} set") report, meter = self.evaluate(getattr(self, f"{dataset_type}_loader"), use_tqdm=True) prefix = f"{report.dataset_name}: full {dataset_type}" self._summarize_report(meter, prefix) def _calculate_time_left(self): time_taken_for_log = time.time() * 1000 - self.train_timer.start iterations_left = self.max_updates - self.num_updates num_logs_left = iterations_left / self.log_interval time_left = num_logs_left * time_taken_for_log snapshot_iteration = self.snapshot_iterations / self.log_interval snapshot_iteration *= iterations_left / self.evaluation_interval time_left += snapshot_iteration * time_taken_for_log return self.train_timer.get_time_hhmmss(gap=time_left) def profile(self, text): if self.not_debug: return self.writer.write(text + ": " + self.profiler.get_time_since_start(), "debug") self.profiler.reset() def predict(self, dataset_type): reporter = self.dataset_loader.get_test_reporter(dataset_type) with torch.no_grad(): self.model.eval() message = f"Starting {dataset_type} inference predictions" self.writer.write(message) while reporter.next_dataset(): dataloader = reporter.get_dataloader() for batch in tqdm(dataloader): prepared_batch = reporter.prepare_batch(batch) model_output = self.model(prepared_batch) report = Report(prepared_batch, model_output) reporter.add_to_report(report, self.model) self.writer.write("Finished predicting") self.model.train()
class TestReporter(Dataset): def __init__(self, multi_task_instance): self.test_task = multi_task_instance self.task_type = multi_task_instance.dataset_type self.config = registry.get("config") self.writer = registry.get("writer") self.report = [] self.timer = Timer() self.training_config = self.config.training self.num_workers = self.training_config.num_workers self.batch_size = self.training_config.batch_size self.report_folder_arg = get_mmf_env(key="report_dir") self.experiment_name = self.training_config.experiment_name self.datasets = [] for dataset in self.test_task.get_datasets(): self.datasets.append(dataset) self.current_dataset_idx = -1 self.current_dataset = self.datasets[self.current_dataset_idx] self.save_dir = get_mmf_env(key="save_dir") self.report_folder = ckpt_name_from_core_args(self.config) self.report_folder += foldername_from_config_override(self.config) self.report_folder = os.path.join(self.save_dir, self.report_folder) self.report_folder = os.path.join(self.report_folder, "reports") if self.report_folder_arg: self.report_folder = self.report_folder_arg PathManager.mkdirs(self.report_folder) def next_dataset(self): if self.current_dataset_idx >= 0: self.flush_report() self.current_dataset_idx += 1 if self.current_dataset_idx == len(self.datasets): return False else: self.current_dataset = self.datasets[self.current_dataset_idx] self.writer.write("Predicting for " + self.current_dataset.dataset_name) return True def flush_report(self): if not is_master(): return name = self.current_dataset.dataset_name time_format = "%Y-%m-%dT%H:%M:%S" time = self.timer.get_time_hhmmss(None, format=time_format) filename = name + "_" if len(self.experiment_name) > 0: filename += self.experiment_name + "_" filename += self.task_type + "_" filename += time + ".json" filepath = os.path.join(self.report_folder, filename) with PathManager.open(filepath, "w") as f: json.dump(self.report, f) self.writer.write( "Wrote evalai predictions for %s to %s" % (name, os.path.abspath(filepath)) ) self.report = [] def get_dataloader(self): other_args = self._add_extra_args_for_dataloader() return DataLoader( dataset=self.current_dataset, collate_fn=BatchCollator( self.current_dataset.dataset_name, self.current_dataset.dataset_type ), num_workers=self.num_workers, pin_memory=self.config.training.pin_memory, **other_args ) def _add_extra_args_for_dataloader(self, other_args=None): if other_args is None: other_args = {} if torch.distributed.is_initialized(): other_args["sampler"] = DistributedSampler( self.current_dataset, shuffle=False ) else: other_args["shuffle"] = False other_args["batch_size"] = get_batch_size() return other_args def prepare_batch(self, batch): return self.current_dataset.prepare_batch(batch) def __len__(self): return len(self.current_dataset) def __getitem__(self, idx): return self.current_dataset[idx] def add_to_report(self, report): # TODO: Later gather whole report for no opinions if self.current_dataset.dataset_name == "coco": report.captions = gather_tensor(report.captions) if isinstance(report.image_id, torch.Tensor): report.image_id = gather_tensor(report.image_id).view(-1) else: report.scores = gather_tensor(report.scores).view( -1, report.scores.size(-1) ) if "question_id" in report: report.question_id = gather_tensor(report.question_id).view(-1) if "image_id" in report: _, enc_size = report.image_id.size() report.image_id = gather_tensor(report.image_id) report.image_id = report.image_id.view(-1, enc_size) if "context_tokens" in report: _, enc_size = report.context_tokens.size() report.context_tokens = gather_tensor(report.context_tokens) report.context_tokens = report.context_tokens.view(-1, enc_size) if not is_master(): return results = self.current_dataset.format_for_evalai(report) self.report = self.report + results
class TestReporter(Dataset): def __init__(self, multi_task_instance): self.test_task = multi_task_instance self.task_type = multi_task_instance.dataset_type self.config = registry.get("config") self.report = [] self.timer = Timer() self.training_config = self.config.training self.num_workers = self.training_config.num_workers self.batch_size = self.training_config.batch_size self.report_folder_arg = get_mmf_env(key="report_dir") self.experiment_name = self.training_config.experiment_name self.datasets = [] for dataset in self.test_task.get_datasets(): self.datasets.append(dataset) self.current_dataset_idx = -1 self.current_dataset = self.datasets[self.current_dataset_idx] self.save_dir = get_mmf_env(key="save_dir") self.report_folder = ckpt_name_from_core_args(self.config) self.report_folder += foldername_from_config_override(self.config) self.report_folder = os.path.join(self.save_dir, self.report_folder) self.report_folder = os.path.join(self.report_folder, "reports") if self.report_folder_arg: self.report_folder = self.report_folder_arg PathManager.mkdirs(self.report_folder) def next_dataset(self): if self.current_dataset_idx >= 0: self.flush_report() self.current_dataset_idx += 1 if self.current_dataset_idx == len(self.datasets): return False else: self.current_dataset = self.datasets[self.current_dataset_idx] logger.info(f"Predicting for {self.current_dataset.dataset_name}") return True def flush_report(self): if not is_master(): return name = self.current_dataset.dataset_name time_format = "%Y-%m-%dT%H:%M:%S" time = self.timer.get_time_hhmmss(None, format=time_format) filename = name + "_" if len(self.experiment_name) > 0: filename += self.experiment_name + "_" filename += self.task_type + "_" filename += time if self.config.evaluation.predict_file_format == "csv": filepath = os.path.join(self.report_folder, filename + ".csv") self.csv_dump(filepath) else: filepath = os.path.join(self.report_folder, filename + ".json") self.json_dump(filepath) logger.info( f"Wrote evalai predictions for {name} to {os.path.abspath(filepath)}" ) self.report = [] def csv_dump(self, filepath): with PathManager.open(filepath, "w") as f: title = self.report[0].keys() cw = csv.DictWriter(f, title, delimiter=",", quoting=csv.QUOTE_MINIMAL) cw.writeheader() cw.writerows(self.report) def json_dump(self, filepath): with PathManager.open(filepath, "w") as f: json.dump(self.report, f) def get_dataloader(self): other_args = self._add_extra_args_for_dataloader() return DataLoader( dataset=self.current_dataset, collate_fn=BatchCollator( self.current_dataset.dataset_name, self.current_dataset.dataset_type ), num_workers=self.num_workers, pin_memory=self.config.training.pin_memory, **other_args, ) def _add_extra_args_for_dataloader(self, other_args=None): if other_args is None: other_args = {} if is_dist_initialized(): other_args["sampler"] = DistributedSampler( self.current_dataset, shuffle=False ) else: other_args["shuffle"] = False other_args["batch_size"] = get_batch_size() return other_args def prepare_batch(self, batch): return self.current_dataset.prepare_batch(batch) def __len__(self): return len(self.current_dataset) def __getitem__(self, idx): return self.current_dataset[idx] def add_to_report(self, report, model): keys = ["id", "question_id", "image_id", "context_tokens", "captions", "scores"] for key in keys: report = self.reshape_and_gather(report, key) if not is_master(): return results = self.current_dataset.format_for_prediction(report) if hasattr(model, "format_for_prediction"): results = model.format_for_prediction(results, report) elif hasattr(model.module, "format_for_prediction"): results = model.module.format_for_prediction(results, report) self.report = self.report + results def reshape_and_gather(self, report, key): if key in report: num_dims = report[key].dim() if num_dims == 1: report[key] = gather_tensor(report[key]).view(-1) elif num_dims >= 2: # Collect dims other than batch other_dims = report[key].size()[1:] report[key] = gather_tensor(report[key]).view(-1, *other_dims) return report
class TestReporter(Dataset): @dataclass class Config: # A set of fields to be *considered* for exporting by the reporter # Note that `format_for_prediction` is what ultimtly detemrimes the # exported fields candidate_fields: List[str] = field( default_factory=lambda: DEFAULT_CANDIDATE_FIELDS) # csv or json predict_file_format: str = "json" def __init__( self, datamodules: List[pl.LightningDataModule], config: Config = None, dataset_type: str = "train", ): self.test_reporter_config = OmegaConf.merge( OmegaConf.structured(self.Config), config) self.datamodules = datamodules self.dataset_type = dataset_type self.config = registry.get("config") self.report = [] self.timer = Timer() self.training_config = self.config.training self.num_workers = self.training_config.num_workers self.batch_size = self.training_config.batch_size self.report_folder_arg = get_mmf_env(key="report_dir") self.experiment_name = self.training_config.experiment_name self.current_datamodule_idx = -1 self.dataset_names = list(self.datamodules.keys()) self.current_datamodule = self.datamodules[self.dataset_names[ self.current_datamodule_idx]] self.current_dataloader = None self.save_dir = get_mmf_env(key="save_dir") self.report_folder = ckpt_name_from_core_args(self.config) self.report_folder += foldername_from_config_override(self.config) self.report_folder = os.path.join(self.save_dir, self.report_folder) self.report_folder = os.path.join(self.report_folder, "reports") if self.report_folder_arg: self.report_folder = self.report_folder_arg self.candidate_fields = self.test_reporter_config.candidate_fields PathManager.mkdirs(self.report_folder) @property def current_dataset(self): self._check_current_dataloader() return self.current_dataloader.dataset def next_dataset(self, flush_report=True): if self.current_datamodule_idx >= 0: if flush_report: self.flush_report() else: self.report = [] self.current_datamodule_idx += 1 if self.current_datamodule_idx == len(self.datamodules): return False else: self.current_datamodule = self.datamodules[self.dataset_names[ self.current_datamodule_idx]] logger.info( f"Predicting for {self.dataset_names[self.current_datamodule_idx]}" ) return True def flush_report(self): if not is_master(): return name = self.current_datamodule.dataset_name time_format = "%Y-%m-%dT%H:%M:%S" time = self.timer.get_time_hhmmss(None, format=time_format) filename = name + "_" if len(self.experiment_name) > 0: filename += self.experiment_name + "_" filename += self.dataset_type + "_" filename += time use_csv_writer = (self.config.evaluation.predict_file_format == "csv" or self.test_reporter_config.predict_file_format == "csv") if use_csv_writer: filepath = os.path.join(self.report_folder, filename + ".csv") self.csv_dump(filepath) else: filepath = os.path.join(self.report_folder, filename + ".json") self.json_dump(filepath) logger.info( f"Wrote predictions for {name} to {os.path.abspath(filepath)}") self.report = [] def postprocess_dataset_report(self): self._check_current_dataloader() if hasattr(self.current_dataset, "on_prediction_end"): self.report = self.current_dataset.on_prediction_end(self.report) def csv_dump(self, filepath): with PathManager.open(filepath, "w") as f: title = self.report[0].keys() cw = csv.DictWriter(f, title, delimiter=",", quoting=csv.QUOTE_MINIMAL) cw.writeheader() cw.writerows(self.report) def json_dump(self, filepath): with PathManager.open(filepath, "w") as f: json.dump(self.report, f) def get_dataloader(self): self.current_dataloader = getattr(self.current_datamodule, f"{self.dataset_type}_dataloader")() # Make sure to assign dataset to dataloader object as # required by MMF if not hasattr(self.current_dataloader, "dataset"): self.current_dataloader.dataset = getattr( self.current_datamodule, f"{self.dataset_type}_dataset") return self.current_dataloader def prepare_batch(self, batch): self._check_current_dataloader() if hasattr(self.current_dataset, "prepare_batch"): batch = self.current_dataset.prepare_batch(batch) batch = convert_batch_to_sample_list(batch) batch.dataset_name = self.current_dataset.dataset_name batch.dataset_type = self.dataset_type return batch def __len__(self): self._check_current_dataloader() return len(self.current_dataloader) def _check_current_dataloader(self): assert self.current_dataloader is not None, ( "Please call `get_dataloader` before accessing any " + "'current_dataloader' based function") def add_to_report(self, report, model, *args, **kwargs): if "execute_on_master_only" in kwargs: warnings.warn( "'execute_on_master_only keyword is deprecated and isn't used anymore", DeprecationWarning, ) self._check_current_dataloader() for key in self.candidate_fields: report = self.reshape_and_gather(report, key) results = [] if hasattr(self.current_dataset, "format_for_prediction"): results = self.current_dataset.format_for_prediction(report) if hasattr(model, "format_for_prediction"): results = model.format_for_prediction(results, report) elif hasattr(model.module, "format_for_prediction"): results = model.module.format_for_prediction(results, report) self.report = self.report + results def reshape_and_gather(self, report, key): if key in report: num_dims = report[key].dim() if num_dims == 1: report[key] = gather_tensor(report[key]).view(-1) elif num_dims >= 2: # Collect dims other than batch other_dims = report[key].size()[1:] report[key] = gather_tensor(report[key]).view(-1, *other_dims) return report
class Logger: def __init__(self, config, name=None): self.logger = None self._is_master = is_master() self.timer = Timer() self.config = config self.save_dir = get_mmf_env(key="save_dir") self.log_format = config.training.log_format self.time_format = "%Y-%m-%dT%H:%M:%S" self.log_filename = "train_" self.log_filename += self.timer.get_time_hhmmss(None, format=self.time_format) self.log_filename += ".log" self.log_folder = os.path.join(self.save_dir, "logs") env_log_dir = get_mmf_env(key="log_dir") if env_log_dir: self.log_folder = env_log_dir if not PathManager.exists(self.log_folder): PathManager.mkdirs(self.log_folder) self.log_filename = os.path.join(self.log_folder, self.log_filename) if not self._is_master: return if self._is_master: print("Logging to:", self.log_filename) logging.captureWarnings(True) if not name: name = __name__ self.logger = logging.getLogger(name) self._file_only_logger = logging.getLogger(name) warnings_logger = logging.getLogger("py.warnings") # Set level level = config.training.logger_level self.logger.setLevel(getattr(logging, level.upper())) self._file_only_logger.setLevel(getattr(logging, level.upper())) formatter = logging.Formatter( "%(asctime)s %(levelname)s: %(message)s", datefmt="%Y-%m-%dT%H:%M:%S" ) # Add handler to file channel = logging.FileHandler(filename=self.log_filename, mode="a") channel.setFormatter(formatter) self.logger.addHandler(channel) self._file_only_logger.addHandler(channel) warnings_logger.addHandler(channel) # Add handler to stdout channel = logging.StreamHandler(sys.stdout) channel.setFormatter(formatter) self.logger.addHandler(channel) warnings_logger.addHandler(channel) should_not_log = self.config.training.should_not_log self.should_log = not should_not_log # Single log wrapper map self._single_log_map = set() def write(self, x, level="info", donot_print=False, log_all=False): if self.logger is None: return if log_all is False and not self._is_master: return # if it should not log then just print it if self.should_log: if hasattr(self.logger, level): if donot_print: getattr(self._file_only_logger, level)(str(x)) else: getattr(self.logger, level)(str(x)) else: self.logger.error("Unknown log level type: %s" % level) else: print(str(x) + "\n") def log_progress(self, info): if not isinstance(info, collections.Mapping): self.write(info) if not self._is_master: return if self.log_format == "simple": output = ", ".join([f"{key}: {value}" for key, value in info.items()]) elif self.log_format == "json": output = json.dumps(info) else: output = str(info) self.write(output) def single_write(self, x, level="info", log_all=False): if self.logger is None: return if log_all is False and not self._is_master: return if x + "_" + level in self._single_log_map: return else: self.write(x, level)
class TestReporter(Dataset): @dataclass class TestReporterConfigType: # A set of fields to be *considered* for exporting by the reporter # Note that `format_for_prediction` is what ultimtly detemrimes the # exported fields candidate_fields: List[str] = MISSING # csv or json predict_file_format: str = MISSING def __init__(self, multi_task_instance, test_reporter_config: TestReporterConfigType = None): if not isinstance(test_reporter_config, TestReporter.TestReporterConfigType): test_reporter_config = TestReporter.TestReporterConfigType( **test_reporter_config) self.test_task = multi_task_instance self.task_type = multi_task_instance.dataset_type self.config = registry.get("config") self.report = [] self.timer = Timer() self.training_config = self.config.training self.num_workers = self.training_config.num_workers self.batch_size = self.training_config.batch_size self.report_folder_arg = get_mmf_env(key="report_dir") self.experiment_name = self.training_config.experiment_name self.test_reporter_config = test_reporter_config self.datasets = [] for dataset in self.test_task.get_datasets(): self.datasets.append(dataset) self.current_dataset_idx = -1 self.current_dataset = self.datasets[self.current_dataset_idx] self.save_dir = get_mmf_env(key="save_dir") self.report_folder = ckpt_name_from_core_args(self.config) self.report_folder += foldername_from_config_override(self.config) self.report_folder = os.path.join(self.save_dir, self.report_folder) self.report_folder = os.path.join(self.report_folder, "reports") if self.report_folder_arg: self.report_folder = self.report_folder_arg self.candidate_fields = DEFAULT_CANDIDATE_FIELDS if not test_reporter_config.candidate_fields == MISSING: self.candidate_fields = test_reporter_config.candidate_fields PathManager.mkdirs(self.report_folder) def next_dataset(self, flush_report=True): if self.current_dataset_idx >= 0: if flush_report: self.flush_report() else: self.report = [] self.current_dataset_idx += 1 if self.current_dataset_idx == len(self.datasets): return False else: self.current_dataset = self.datasets[self.current_dataset_idx] logger.info(f"Predicting for {self.current_dataset.dataset_name}") return True def flush_report(self): if not is_master(): return name = self.current_dataset.dataset_name time_format = "%Y-%m-%dT%H:%M:%S" time = self.timer.get_time_hhmmss(None, format=time_format) filename = name + "_" if len(self.experiment_name) > 0: filename += self.experiment_name + "_" filename += self.task_type + "_" filename += time use_csv_writer = (self.config.evaluation.predict_file_format == "csv" or self.test_reporter_config.predict_file_format == "csv") if use_csv_writer: filepath = os.path.join(self.report_folder, filename + ".csv") self.csv_dump(filepath) else: filepath = os.path.join(self.report_folder, filename + ".json") self.json_dump(filepath) logger.info( f"Wrote predictions for {name} to {os.path.abspath(filepath)}") self.report = [] def postprocess_dataset_report(self): if hasattr(self.current_dataset, "on_prediction_end"): self.report = self.current_dataset.on_prediction_end(self.report) def csv_dump(self, filepath): with PathManager.open(filepath, "w") as f: title = self.report[0].keys() cw = csv.DictWriter(f, title, delimiter=",", quoting=csv.QUOTE_MINIMAL) cw.writeheader() cw.writerows(self.report) def json_dump(self, filepath): with PathManager.open(filepath, "w") as f: json.dump(self.report, f) def get_dataloader(self): dataloader, _ = build_dataloader_and_sampler(self.current_dataset, self.training_config) return dataloader def prepare_batch(self, batch): if hasattr(self.current_dataset, "prepare_batch"): batch = self.current_dataset.prepare_batch(batch) return batch def __len__(self): return len(self.current_dataset) def __getitem__(self, idx): return self.current_dataset[idx] def add_to_report(self, report, model, execute_on_master_only=True): for key in self.candidate_fields: report = self.reshape_and_gather(report, key) if execute_on_master_only and not is_master(): return results = [] if hasattr(self.current_dataset, "format_for_prediction"): results = self.current_dataset.format_for_prediction(report) if hasattr(model, "format_for_prediction"): results = model.format_for_prediction(results, report) elif hasattr(model.module, "format_for_prediction"): results = model.module.format_for_prediction(results, report) self.report = self.report + results def reshape_and_gather(self, report, key): if key in report: num_dims = report[key].dim() if num_dims == 1: report[key] = gather_tensor(report[key]).view(-1) elif num_dims >= 2: # Collect dims other than batch other_dims = report[key].size()[1:] report[key] = gather_tensor(report[key]).view(-1, *other_dims) return report