| | |
| | """ |
| | Training script using the new "LazyConfig" python config files. |
| | |
| | This scripts reads a given python config file and runs the training or evaluation. |
| | It can be used to train any models or dataset as long as they can be |
| | instantiated by the recursive construction defined in the given config file. |
| | |
| | Besides lazy construction of models, dataloader, etc., this scripts expects a |
| | few common configuration parameters currently defined in "configs/common/train.py". |
| | To add more complicated training logic, you can easily add other configs |
| | in the config file and implement a new train_net.py to handle them. |
| | """ |
| | import logging |
| | import os |
| | import random |
| | import sys |
| | import time |
| | from collections import abc |
| | from contextlib import nullcontext |
| | from datetime import timedelta |
| |
|
| | import torch |
| | from torch.nn.parallel import DataParallel, DistributedDataParallel |
| |
|
| | import ape |
| | from ape.checkpoint import DetectionCheckpointer |
| | from ape.engine import SimpleTrainer |
| | from ape.evaluation import inference_on_dataset |
| | from detectron2.config import LazyConfig, instantiate |
| | from detectron2.engine import default_argument_parser |
| | from detectron2.engine import default_setup, hooks, launch |
| | from detectron2.engine.defaults import create_ddp_model |
| | from detectron2.evaluation import print_csv_format |
| | from detectron2.utils import comm |
| | from detectron2.utils.events import ( |
| | CommonMetricPrinter, |
| | JSONWriter, |
| | TensorboardXWriter, |
| | get_event_storage, |
| | ) |
| | from detectron2.utils.file_io import PathManager |
| | from detectron2.utils.logger import setup_logger |
| | from detrex.modeling import ema |
| | from detrex.utils import WandbWriter |
| |
|
| | sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), os.path.pardir))) |
| |
|
| | logger = logging.getLogger("ape") |
| |
|
| |
|
| | class Trainer(SimpleTrainer): |
| | """ |
| | We've combine Simple and AMP Trainer together. |
| | """ |
| |
|
| | def __init__( |
| | self, |
| | model, |
| | dataloader, |
| | optimizer, |
| | amp=False, |
| | clip_grad_params=None, |
| | grad_scaler=None, |
| | iter_size=1, |
| | iter_loop=True, |
| | dataset_ratio=None, |
| | save_memory=False, |
| | ): |
| | super().__init__(model=model, data_loader=dataloader, optimizer=optimizer) |
| |
|
| | unsupported = "AMPTrainer does not support single-process multi-device training!" |
| | if isinstance(model, DistributedDataParallel): |
| | assert not (model.device_ids and len(model.device_ids) > 1), unsupported |
| | assert not isinstance(model, DataParallel), unsupported |
| |
|
| | if amp: |
| | if grad_scaler is None: |
| | from torch.cuda.amp import GradScaler |
| |
|
| | grad_scaler = GradScaler() |
| | self.grad_scaler = grad_scaler |
| |
|
| | self.amp = amp |
| |
|
| | self.clip_grad_params = clip_grad_params |
| |
|
| | if isinstance(model, DistributedDataParallel): |
| | if hasattr(model.module, "model_vision"): |
| | self.dataset_names = model.module.model_vision.dataset_names |
| | else: |
| | self.dataset_names = ["unknown"] |
| | else: |
| | if hasattr(model, "model_vision"): |
| | self.dataset_names = model.model_vision.dataset_names |
| | else: |
| | self.dataset_names = ["unknown"] |
| | self.dataset_image_counts = { |
| | k: torch.tensor(0, dtype=torch.float).to(comm.get_local_rank()) |
| | for k in self.dataset_names |
| | } |
| | self.dataset_object_counts = { |
| | k: torch.tensor(0, dtype=torch.float).to(comm.get_local_rank()) |
| | for k in self.dataset_names |
| | } |
| |
|
| | self.iter_size = iter_size |
| | self.iter_loop = iter_loop |
| | self.dataset_ratio = dataset_ratio |
| | self.save_memory = save_memory |
| |
|
| | def run_step(self): |
| | if self.iter_size > 1: |
| | if self.iter_loop: |
| | return self.run_step_accumulate_iter_loop() |
| | else: |
| | return self.run_step_accumulate() |
| | """ |
| | Implement the standard training logic described above. |
| | """ |
| | assert self.model.training, "[Trainer] model was changed to eval mode!" |
| | assert torch.cuda.is_available(), "[Trainer] CUDA is required for AMP training!" |
| | from torch.cuda.amp import autocast |
| |
|
| | start = time.perf_counter() |
| | """ |
| | If you want to do something with the data, you can wrap the dataloader. |
| | """ |
| | while True: |
| | data = next(self._data_loader_iter) |
| | if all([len(x["instances"]) > 0 for x in data]): |
| | break |
| | data_time = time.perf_counter() - start |
| |
|
| | for d in data: |
| | if d.get("dataloader_id", None) is not None: |
| | d["dataset_id"] = d["dataloader_id"] |
| | self.dataset_image_counts[self.dataset_names[d.get("dataset_id", 0)]] += 1 |
| | self.dataset_object_counts[self.dataset_names[d.get("dataset_id", 0)]] += len( |
| | d.get("instances", []) |
| | ) |
| | dataset_image_counts = {f"count_image/{k}": v for k, v in self.dataset_image_counts.items()} |
| | dataset_object_counts = { |
| | f"count_object/{k}": v for k, v in self.dataset_object_counts.items() |
| | } |
| | if self.async_write_metrics: |
| | self.concurrent_executor.submit( |
| | self._write_metrics_common, dataset_image_counts, iter=self.iter |
| | ) |
| | self.concurrent_executor.submit( |
| | self._write_metrics_common, dataset_object_counts, iter=self.iter |
| | ) |
| | else: |
| | self._write_metrics_common(dataset_image_counts) |
| | self._write_metrics_common(dataset_object_counts) |
| |
|
| | """ |
| | If you want to do something with the losses, you can wrap the model. |
| | """ |
| | with autocast(enabled=self.amp): |
| | loss_dict = self.model(data) |
| | if isinstance(loss_dict, torch.Tensor): |
| | losses = loss_dict |
| | loss_dict = {"total_loss": loss_dict} |
| | else: |
| | losses = sum(loss_dict.values()) |
| |
|
| | """ |
| | If you need to accumulate gradients or do something similar, you can |
| | wrap the optimizer with your custom `zero_grad()` method. |
| | """ |
| | self.optimizer.zero_grad() |
| |
|
| | if self.amp: |
| | self.grad_scaler.scale(losses).backward() |
| | if self.clip_grad_params is not None: |
| | self.grad_scaler.unscale_(self.optimizer) |
| | self.clip_grads(self.model.parameters()) |
| | self.grad_scaler.step(self.optimizer) |
| | self.grad_scaler.update() |
| | else: |
| | losses.backward() |
| | if self.clip_grad_params is not None: |
| | self.clip_grads(self.model.parameters()) |
| | self.optimizer.step() |
| |
|
| | if self.async_write_metrics: |
| | self.concurrent_executor.submit( |
| | self._write_metrics, loss_dict, data_time, iter=self.iter |
| | ) |
| | else: |
| | self._write_metrics(loss_dict, data_time) |
| |
|
| | if self.save_memory: |
| | del losses |
| | del loss_dict |
| | torch.cuda.empty_cache() |
| |
|
| | def run_step_accumulate(self): |
| | """ |
| | Implement the standard training logic described above. |
| | """ |
| | assert self.model.training, "[Trainer] model was changed to eval mode!" |
| | assert torch.cuda.is_available(), "[Trainer] CUDA is required for AMP training!" |
| | from torch.cuda.amp import autocast |
| |
|
| | start = time.perf_counter() |
| | """ |
| | If you want to do something with the data, you can wrap the dataloader. |
| | """ |
| | while True: |
| | data = next(self._data_loader_iter) |
| | if all([len(x["instances"]) > 0 for x in data]): |
| | break |
| | data_time = time.perf_counter() - start |
| |
|
| | for d in data: |
| | if d.get("dataloader_id", None) is not None: |
| | d["dataset_id"] = d["dataloader_id"] |
| | self.dataset_image_counts[self.dataset_names[d.get("dataset_id", 0)]] += 1 |
| | self.dataset_object_counts[self.dataset_names[d.get("dataset_id", 0)]] += len( |
| | d.get("instances", []) |
| | ) |
| | dataset_image_counts = {f"count_image/{k}": v for k, v in self.dataset_image_counts.items()} |
| | dataset_object_counts = { |
| | f"count_object/{k}": v for k, v in self.dataset_object_counts.items() |
| | } |
| | if self.async_write_metrics: |
| | self.concurrent_executor.submit( |
| | self._write_metrics_common, dataset_image_counts, iter=self.iter |
| | ) |
| | self.concurrent_executor.submit( |
| | self._write_metrics_common, dataset_object_counts, iter=self.iter |
| | ) |
| | else: |
| | self._write_metrics_common(dataset_image_counts) |
| | self._write_metrics_common(dataset_object_counts) |
| |
|
| | sync_context = self.model.no_sync if (self.iter + 1) % self.iter_size != 0 else nullcontext |
| | """ |
| | If you want to do something with the losses, you can wrap the model. |
| | """ |
| | with sync_context(): |
| | with autocast(enabled=self.amp): |
| | loss_dict = self.model(data) |
| |
|
| | if isinstance(loss_dict, torch.Tensor): |
| | losses = loss_dict |
| | loss_dict = {"total_loss": loss_dict} |
| | else: |
| | losses = sum(loss_dict.values()) |
| |
|
| | """ |
| | If you need to accumulate gradients or do something similar, you can |
| | wrap the optimizer with your custom `zero_grad()` method. |
| | """ |
| | if self.iter == self.start_iter: |
| | self.optimizer.zero_grad() |
| |
|
| | if self.iter_size > 1: |
| | losses = losses / self.iter_size |
| |
|
| | if self.amp: |
| | self.grad_scaler.scale(losses).backward() |
| | if (self.iter + 1) % self.iter_size == 0: |
| | if self.clip_grad_params is not None: |
| | self.grad_scaler.unscale_(self.optimizer) |
| | self.clip_grads(self.model.parameters()) |
| | self.grad_scaler.step(self.optimizer) |
| | self.grad_scaler.update() |
| | self.optimizer.zero_grad() |
| | else: |
| | losses.backward() |
| | if (self.iter + 1) % self.iter_size == 0: |
| | if self.clip_grad_params is not None: |
| | self.clip_grads(self.model.parameters()) |
| | self.optimizer.step() |
| | self.optimizer.zero_grad() |
| |
|
| | if self.async_write_metrics: |
| | self.concurrent_executor.submit( |
| | self._write_metrics, loss_dict, data_time, iter=self.iter |
| | ) |
| | else: |
| | self._write_metrics(loss_dict, data_time) |
| |
|
| | if self.save_memory: |
| | del losses |
| | del loss_dict |
| | torch.cuda.empty_cache() |
| |
|
| | def run_step_accumulate_iter_loop(self): |
| | """ |
| | Implement the standard training logic described above. |
| | """ |
| | assert self.model.training, "[Trainer] model was changed to eval mode!" |
| | assert torch.cuda.is_available(), "[Trainer] CUDA is required for AMP training!" |
| | from torch.cuda.amp import autocast |
| |
|
| | self.optimizer.zero_grad() |
| | for inner_iter in range(self.iter_size): |
| | start = time.perf_counter() |
| | """ |
| | If you want to do something with the data, you can wrap the dataloader. |
| | """ |
| | while True: |
| | data = next(self._data_loader_iter) |
| | if all([len(x["instances"]) > 0 for x in data]): |
| | break |
| | data_time = time.perf_counter() - start |
| |
|
| | for d in data: |
| | if d.get("dataloader_id", None) is not None: |
| | d["dataset_id"] = d["dataloader_id"] |
| | self.dataset_image_counts[self.dataset_names[d.get("dataset_id", 0)]] += 1 |
| | self.dataset_object_counts[self.dataset_names[d.get("dataset_id", 0)]] += len( |
| | d.get("instances", []) |
| | ) |
| | dataset_image_counts = { |
| | f"count_image/{k}": v for k, v in self.dataset_image_counts.items() |
| | } |
| | dataset_object_counts = { |
| | f"count_object/{k}": v for k, v in self.dataset_object_counts.items() |
| | } |
| | if self.async_write_metrics: |
| | self.concurrent_executor.submit( |
| | self._write_metrics_common, dataset_image_counts, iter=self.iter |
| | ) |
| | self.concurrent_executor.submit( |
| | self._write_metrics_common, dataset_object_counts, iter=self.iter |
| | ) |
| | else: |
| | self._write_metrics_common(dataset_image_counts) |
| | self._write_metrics_common(dataset_object_counts) |
| |
|
| | sync_context = self.model.no_sync if inner_iter != self.iter_size - 1 else nullcontext |
| | """ |
| | If you want to do something with the losses, you can wrap the model. |
| | """ |
| | with sync_context(): |
| | with autocast(enabled=self.amp): |
| | loss_dict = self.model(data) |
| |
|
| | if isinstance(loss_dict, torch.Tensor): |
| | losses = loss_dict |
| | loss_dict = {"total_loss": loss_dict} |
| | else: |
| | losses = sum(loss_dict.values()) |
| |
|
| | """ |
| | If you need to accumulate gradients or do something similar, you can |
| | wrap the optimizer with your custom `zero_grad()` method. |
| | """ |
| |
|
| | losses = losses / self.iter_size |
| |
|
| | if self.amp: |
| | self.grad_scaler.scale(losses).backward() |
| | else: |
| | losses.backward() |
| |
|
| | if self.async_write_metrics: |
| | self.concurrent_executor.submit( |
| | self._write_metrics, loss_dict, data_time, iter=self.iter |
| | ) |
| | else: |
| | self._write_metrics(loss_dict, data_time) |
| |
|
| | if self.save_memory: |
| | del losses |
| | del loss_dict |
| | torch.cuda.empty_cache() |
| |
|
| | if self.amp: |
| | if self.clip_grad_params is not None: |
| | self.grad_scaler.unscale_(self.optimizer) |
| | self.clip_grads(self.model.parameters()) |
| | self.grad_scaler.step(self.optimizer) |
| | self.grad_scaler.update() |
| | else: |
| | if self.clip_grad_params is not None: |
| | self.clip_grads(self.model.parameters()) |
| | self.optimizer.step() |
| |
|
| | def clip_grads(self, params): |
| | params = list(filter(lambda p: p.requires_grad and p.grad is not None, params)) |
| | if len(params) > 0: |
| | return torch.nn.utils.clip_grad_norm_( |
| | parameters=params, |
| | **self.clip_grad_params, |
| | ) |
| |
|
| | def state_dict(self): |
| | ret = super().state_dict() |
| | if self.grad_scaler and self.amp: |
| | ret["grad_scaler"] = self.grad_scaler.state_dict() |
| | return ret |
| |
|
| | def load_state_dict(self, state_dict): |
| | super().load_state_dict(state_dict) |
| | if self.grad_scaler and self.amp: |
| | self.grad_scaler.load_state_dict(state_dict["grad_scaler"]) |
| |
|
| | @property |
| | def _data_loader_iter(self): |
| | if isinstance(self.data_loader, abc.MutableSequence): |
| | if self._data_loader_iter_obj is None: |
| | self._data_loader_iter_obj = [iter(x) for x in self.data_loader] |
| | self._data_loader_indices = [] |
| |
|
| | if len(self._data_loader_indices) == 0: |
| | self._data_loader_indices = random.choices( |
| | list(range(len(self.data_loader))), weights=self.dataset_ratio, k=10000 |
| | ) |
| | idx = self._data_loader_indices.pop() |
| | return self._data_loader_iter_obj[idx] |
| |
|
| | if self._data_loader_iter_obj is None: |
| | self._data_loader_iter_obj = iter(self.data_loader) |
| | return self._data_loader_iter_obj |
| |
|
| |
|
| | def do_test(cfg, model, eval_only=False): |
| | logger = logging.getLogger("ape") |
| | if "evaluator" in cfg.dataloader: |
| | if isinstance(model, DistributedDataParallel): |
| | if hasattr(model.module, "set_eval_dataset"): |
| | model.module.set_eval_dataset(cfg.dataloader.test.dataset.names) |
| | else: |
| | if hasattr(model, "set_eval_dataset"): |
| | model.set_eval_dataset(cfg.dataloader.test.dataset.names) |
| | output_dir = os.path.join( |
| | cfg.train.output_dir, "inference_{}".format(cfg.dataloader.test.dataset.names) |
| | ) |
| | if "cityscapes" in cfg.dataloader.test.dataset.names: |
| | pass |
| | else: |
| | if isinstance(cfg.dataloader.evaluator, abc.MutableSequence): |
| | for evaluator in cfg.dataloader.evaluator: |
| | evaluator.output_dir = output_dir |
| | else: |
| | cfg.dataloader.evaluator.output_dir = output_dir |
| |
|
| | ret = inference_on_dataset( |
| | model, instantiate(cfg.dataloader.test), instantiate(cfg.dataloader.evaluator) |
| | ) |
| | logger.info( |
| | "Evaluation results for {} in csv format:".format(cfg.dataloader.test.dataset.names) |
| | ) |
| | print_csv_format(ret) |
| | ret = {f"{k}_{cfg.dataloader.test.dataset.names}": v for k, v in ret.items()} |
| | else: |
| | ret = {} |
| |
|
| | if "evaluators" in cfg.dataloader: |
| | for test, evaluator in zip(cfg.dataloader.tests, cfg.dataloader.evaluators): |
| | if isinstance(model, DistributedDataParallel): |
| | model.module.set_eval_dataset(test.dataset.names) |
| | else: |
| | model.set_eval_dataset(test.dataset.names) |
| | output_dir = os.path.join( |
| | cfg.train.output_dir, "inference_{}".format(test.dataset.names) |
| | ) |
| | if isinstance(evaluator, abc.MutableSequence): |
| | for eva in evaluator: |
| | eva.output_dir = output_dir |
| | else: |
| | evaluator.output_dir = output_dir |
| | ret_ = inference_on_dataset(model, instantiate(test), instantiate(evaluator)) |
| | logger.info("Evaluation results for {} in csv format:".format(test.dataset.names)) |
| | print_csv_format(ret_) |
| | ret.update({f"{k}_{test.dataset.names}": v for k, v in ret_.items()}) |
| |
|
| | bbox_odinw_AP = {"AP": [], "AP50": [], "AP75": [], "APs": [], "APm": [], "APl": []} |
| | segm_seginw_AP = {"AP": [], "AP50": [], "AP75": [], "APs": [], "APm": [], "APl": []} |
| | bbox_rf100_AP = {"AP": [], "AP50": [], "AP75": [], "APs": [], "APm": [], "APl": []} |
| | for k, v in ret.items(): |
| | for kk, vv in v.items(): |
| | if k.startswith("bbox_odinw") and kk in bbox_odinw_AP and vv == vv: |
| | bbox_odinw_AP[kk].append(vv) |
| | if k.startswith("segm_seginw") and kk in segm_seginw_AP and vv == vv: |
| | segm_seginw_AP[kk].append(vv) |
| | if k.startswith("bbox_rf100") and kk in bbox_rf100_AP and vv == vv: |
| | bbox_rf100_AP[kk].append(vv) |
| |
|
| | from statistics import median, mean |
| |
|
| | logger.info("Evaluation results: {}".format(ret)) |
| | for k, v in bbox_odinw_AP.items(): |
| | if len(v) > 0: |
| | logger.info( |
| | "Evaluation results for odinw bbox {}: mean {} median {}".format( |
| | k, mean(v), median(v) |
| | ) |
| | ) |
| | for k, v in segm_seginw_AP.items(): |
| | if len(v) > 0: |
| | logger.info( |
| | "Evaluation results for seginw segm {}: mean {} median {}".format( |
| | k, mean(v), median(v) |
| | ) |
| | ) |
| | for k, v in bbox_rf100_AP.items(): |
| | if len(v) > 0: |
| | logger.info( |
| | "Evaluation results for rf100 bbox {}: mean {} median {}".format( |
| | k, mean(v), median(v) |
| | ) |
| | ) |
| |
|
| | return ret |
| |
|
| |
|
| | def do_train(args, cfg): |
| | """ |
| | Args: |
| | cfg: an object with the following attributes: |
| | model: instantiate to a module |
| | dataloader.{train,test}: instantiate to dataloaders |
| | dataloader.evaluator: instantiate to evaluator for test set |
| | optimizer: instantaite to an optimizer |
| | lr_multiplier: instantiate to a fvcore scheduler |
| | train: other misc config defined in `configs/common/train.py`, including: |
| | output_dir (str) |
| | init_checkpoint (str) |
| | amp.enabled (bool) |
| | max_iter (int) |
| | eval_period, log_period (int) |
| | device (str) |
| | checkpointer (dict) |
| | ddp (dict) |
| | """ |
| | model = instantiate(cfg.model) |
| | logger = logging.getLogger("ape") |
| | logger.info("Model:\n{}".format(model)) |
| | model.to(cfg.train.device) |
| |
|
| | cfg.optimizer.params.model = model |
| | optim = instantiate(cfg.optimizer) |
| |
|
| | if "wait_group" in cfg.dataloader: |
| | wait = comm.get_local_rank() % cfg.dataloader.wait_group * cfg.dataloader.wait_time |
| | logger.info("rank {} sleep {}".format(comm.get_local_rank(), wait)) |
| | time.sleep(wait) |
| | if isinstance(cfg.dataloader.train, abc.MutableSequence): |
| | train_loader = [instantiate(x) for x in cfg.dataloader.train] |
| | else: |
| | train_loader = instantiate(cfg.dataloader.train) |
| |
|
| | model = create_ddp_model(model, **cfg.train.ddp) |
| |
|
| | ema.may_build_model_ema(cfg, model) |
| |
|
| | trainer = Trainer( |
| | model=model, |
| | dataloader=train_loader, |
| | optimizer=optim, |
| | amp=cfg.train.amp.enabled, |
| | clip_grad_params=cfg.train.clip_grad.params if cfg.train.clip_grad.enabled else None, |
| | iter_size=cfg.train.iter_size if "iter_size" in cfg.train else 1, |
| | iter_loop=cfg.train.iter_loop if "iter_loop" in cfg.train else True, |
| | dataset_ratio=cfg.train.dataset_ratio if "dataset_ratio" in cfg.train else None, |
| | ) |
| |
|
| | checkpointer = DetectionCheckpointer( |
| | model, |
| | cfg.train.output_dir, |
| | trainer=trainer, |
| | **ema.may_get_ema_checkpointer(cfg, model), |
| | ) |
| |
|
| | if comm.is_main_process(): |
| | output_dir = cfg.train.output_dir |
| | PathManager.mkdirs(output_dir) |
| | writers = [ |
| | CommonMetricPrinter(cfg.train.max_iter), |
| | JSONWriter(os.path.join(output_dir, "metrics.json")), |
| | TensorboardXWriter(output_dir), |
| | ] |
| | if cfg.train.wandb.enabled: |
| | PathManager.mkdirs(cfg.train.wandb.params.dir) |
| | writers.append(WandbWriter(cfg)) |
| |
|
| | trainer.register_hooks( |
| | [ |
| | hooks.IterationTimer(), |
| | ema.EMAHook(cfg, model) if cfg.train.model_ema.enabled else None, |
| | hooks.LRScheduler(scheduler=instantiate(cfg.lr_multiplier)), |
| | hooks.PeriodicCheckpointer(checkpointer, **cfg.train.checkpointer) |
| | if comm.is_main_process() |
| | else None, |
| | hooks.EvalHook(cfg.train.eval_period, lambda: do_test(cfg, model)), |
| | hooks.PeriodicWriter( |
| | writers, |
| | period=cfg.train.log_period, |
| | ) |
| | if comm.is_main_process() |
| | else None, |
| | ] |
| | ) |
| |
|
| | checkpointer.resume_or_load(cfg.train.init_checkpoint, resume=args.resume) |
| | if args.resume and checkpointer.has_checkpoint(): |
| | start_iter = trainer.iter + 1 |
| | else: |
| | start_iter = 0 |
| | trainer.train(start_iter, cfg.train.max_iter) |
| |
|
| |
|
| | def main(args): |
| | cfg = LazyConfig.load(args.config_file) |
| | cfg = LazyConfig.apply_overrides(cfg, args.opts) |
| |
|
| | if "output_dir" in cfg.model: |
| | cfg.model.output_dir = cfg.train.output_dir |
| | if "model_vision" in cfg.model and "output_dir" in cfg.model.model_vision: |
| | cfg.model.model_vision.output_dir = cfg.train.output_dir |
| | if "train" in cfg.dataloader: |
| | if isinstance(cfg.dataloader.train, abc.MutableSequence): |
| | for i in range(len(cfg.dataloader.train)): |
| | if "output_dir" in cfg.dataloader.train[i].mapper: |
| | cfg.dataloader.train[i].mapper.output_dir = cfg.train.output_dir |
| | else: |
| | if "output_dir" in cfg.dataloader.train.mapper: |
| | cfg.dataloader.train.mapper.output_dir = cfg.train.output_dir |
| |
|
| | default_setup(cfg, args) |
| |
|
| | setup_logger(cfg.train.output_dir, distributed_rank=comm.get_rank(), name="ape") |
| | setup_logger(cfg.train.output_dir, distributed_rank=comm.get_rank(), name="timm") |
| |
|
| | if cfg.train.fast_dev_run.enabled: |
| | cfg.train.max_iter = 20 |
| | cfg.train.eval_period = 10 |
| | cfg.train.log_period = 1 |
| |
|
| | if args.eval_only: |
| | model = instantiate(cfg.model) |
| | logger = logging.getLogger("ape") |
| | logger.info("Model:\n{}".format(model)) |
| | model.to(cfg.train.device) |
| | model.to(torch.float16) |
| | model = create_ddp_model(model) |
| |
|
| | ema.may_build_model_ema(cfg, model) |
| | DetectionCheckpointer(model, **ema.may_get_ema_checkpointer(cfg, model)).load( |
| | cfg.train.init_checkpoint |
| | ) |
| | if cfg.train.model_ema.enabled and cfg.train.model_ema.use_ema_weights_for_eval_only: |
| | ema.apply_model_ema(model) |
| | print(do_test(cfg, model, eval_only=True)) |
| | else: |
| | do_train(args, cfg) |
| |
|
| |
|
| | if __name__ == "__main__": |
| | args = default_argument_parser().parse_args() |
| | launch( |
| | main, |
| | args.num_gpus, |
| | num_machines=args.num_machines, |
| | machine_rank=args.machine_rank, |
| | dist_url=args.dist_url, |
| | args=(args,), |
| | timeout=timedelta(minutes=120), |
| | ) |
| |
|
