def __init__(self, cfg): logger = logging.getLogger("detectron2") if not logger.isEnabledFor(logging.INFO): setup_logger() cfg = DefaultTrainer.auto_scale_workers(cfg, comm.get_world_size()) # Assume these objects must be constructed in this order. model = self.build_model(cfg) optimizer = self.build_optimizer(cfg, model) data_loader, num_per_epoch = self.build_train_loader(cfg) # update iteration cfg to epoch cfg if cfg.SOLVER.EPOCH.ENABLED: cfg = self.adjust_epoch_to_iter(cfg, num_per_epoch) # For training, wrap with DDP. But don't need this for inference. if comm.get_world_size() > 1: model = DistributedDataParallel(model, device_ids=[comm.get_local_rank()], broadcast_buffers=False) super(DefaultTrainer, self).__init__(model, data_loader, optimizer) self.scheduler = self.build_lr_scheduler(cfg, optimizer) # Assume no other objects need to be checkpointed. # We can later make it checkpoint the stateful hooks self.checkpointer = DetectionCheckpointer( # Assume you want to save checkpoints together with logs/statistics model, cfg.OUTPUT_DIR, optimizer=optimizer, scheduler=self.scheduler, ) self.start_iter = 0 self.max_iter = cfg.SOLVER.MAX_ITER self.cfg = cfg self.register_hooks(self.build_hooks())
def __init__(self, cfg): """ Args: cfg (CfgNode): Use the custom checkpointer, which loads other backbone models with matching heuristics. """ cfg = DefaultTrainer.auto_scale_workers(cfg, comm.get_world_size()) model = self.build_model(cfg) optimizer = self.build_optimizer(cfg, model) data_loader = self.build_train_loader(cfg) if comm.get_world_size() > 1: model = DistributedDataParallel(model, device_ids=[comm.get_local_rank()], broadcast_buffers=False) TrainerBase.__init__(self) self._trainer = (AMPTrainer if cfg.SOLVER.AMP.ENABLED else SimpleTrainer)(model, data_loader, optimizer) self.scheduler = self.build_lr_scheduler(cfg, optimizer) self.checkpointer = DetectionCheckpointer( model, cfg.OUTPUT_DIR, optimizer=optimizer, scheduler=self.scheduler, ) self.start_iter = 0 self.max_iter = cfg.SOLVER.MAX_ITER self.cfg = cfg self.register_hooks(self.build_hooks())
def __init__(self, cfg): super().__init__() if not logger.isEnabledFor(logging.INFO): # setup_logger is not called for d2 setup_logger() self.cfg = DefaultTrainer.auto_scale_workers(cfg, comm.get_world_size()) self.storage: EventStorage = None self.model = build_model(self.cfg) self.start_iter = 0 self.max_iter = cfg.SOLVER.MAX_ITER
def __init__(self, cfg): """ Args: cfg (CfgNode): """ logger = logging.getLogger("detectron2") if not logger.isEnabledFor( logging.INFO): # setup_logger is not called for d2 setup_logger() cfg = DefaultTrainer.auto_scale_workers(cfg, comm.get_world_size()) # Assume these objects must be constructed in this order. model = self.build_model(cfg) # add decoupled training setting if cfg.DECOUPLE_TRAINING: logger.info("Decouple Training. Only training the last FC.") for name, param in model.named_parameters(): if 'cls_score' not in name: param.requires_grad = False optimizer = self.build_optimizer(cfg, model) data_loader = self.build_train_loader(cfg) # For training, wrap with DDP. But don't need this for inference. if comm.get_world_size() > 1: model = DistributedDataParallel( model, device_ids=[comm.get_local_rank()], broadcast_buffers=False, ) super(DefaultTrainer, self).__init__(model, data_loader, optimizer) self.scheduler = self.build_lr_scheduler(cfg, optimizer) # Assume no other objects need to be checkpointed. # We can later make it checkpoint the stateful hooks self.checkpointer = DetectionCheckpointer( # Assume you want to save checkpoints together with logs/statistics model, cfg.OUTPUT_DIR, optimizer=optimizer, scheduler=self.scheduler, ) self.start_iter = 0 self.max_iter = cfg.SOLVER.MAX_ITER self.cfg = cfg self.register_hooks(self.build_hooks())
def __init__(self, cfg): """Initializes the CSDTrainer. Most of the code is from `super.__init__()`, the only change is that for `self._trainer` the `CSDTrainer` is used and weight scheduling parameters are injected into it, look for "CSD: ... " comments. """ TrainerBase.__init__(self) # CSD: don't call `super`'s init as we are overriding it logger = logging.getLogger("detectron2") if not logger.isEnabledFor(logging.INFO): # setup_logger is not called for d2 setup_logger() cfg = DefaultTrainer.auto_scale_workers(cfg, comm.get_world_size()) # Assume these objects must be constructed in this order. model = self.build_model(cfg) optimizer = self.build_optimizer(cfg, model) data_loader = self.build_train_loader(cfg) model = create_ddp_model(model, broadcast_buffers=False) self._trainer = CSDTrainer(model, data_loader, optimizer) # CSD: use a CSD-specific trainer # CSD: inject weight scheduling parameters into trainer ( self._trainer.solver_csd_beta, self._trainer.solver_csd_t0, self._trainer.solver_csd_t1, self._trainer.solver_csd_t2, self._trainer.solver_csd_t, ) = ( cfg.SOLVER.CSD_WEIGHT_SCHEDULE_RAMP_BETA, cfg.SOLVER.CSD_WEIGHT_SCHEDULE_RAMP_T0, cfg.SOLVER.CSD_WEIGHT_SCHEDULE_RAMP_T1, cfg.SOLVER.CSD_WEIGHT_SCHEDULE_RAMP_T2, cfg.SOLVER.CSD_WEIGHT_SCHEDULE_RAMP_T, ) self.scheduler = self.build_lr_scheduler(cfg, optimizer) self.checkpointer = WandbDetectionCheckpointer( # CSD: use custom checkpointer (only few lines are added there) # Assume you want to save checkpoints together with logs/statistics model, cfg.OUTPUT_DIR, trainer=weakref.proxy(self), ) self.start_iter = 0 self.max_iter = cfg.SOLVER.MAX_ITER self.cfg = cfg self.register_hooks(self.build_hooks())
def __init__(self, cfg): """ Args: cfg (CfgNode): Use the custom checkpointer, which loads other backbone models with matching heuristics. """ cfg = DefaultTrainer.auto_scale_workers(cfg, comm.get_world_size()) data_loader = self.build_train_loader(cfg) # create an student model model = self.build_model(cfg) optimizer = self.build_optimizer(cfg, model) # create an teacher model model_teacher = self.build_model(cfg) self.model_teacher = model_teacher # For training, wrap with DDP. But don't need this for inference. if comm.get_world_size() > 1: model = DistributedDataParallel( model, device_ids=[comm.get_local_rank()], broadcast_buffers=False ) TrainerBase.__init__(self) self._trainer = (AMPTrainer if cfg.SOLVER.AMP.ENABLED else SimpleTrainer)( model, data_loader, optimizer ) self.scheduler = self.build_lr_scheduler(cfg, optimizer) # Ensemble teacher and student model is for model saving and loading ensem_ts_model = EnsembleTSModel(model_teacher, model) self.checkpointer = DetectionTSCheckpointer( ensem_ts_model, cfg.OUTPUT_DIR, optimizer=optimizer, scheduler=self.scheduler, ) self.start_iter = 0 self.max_iter = cfg.SOLVER.MAX_ITER self.cfg = cfg self.register_hooks(self.build_hooks())
def __init__(self, cfg): super().__init__() self.cfg = DefaultTrainer.auto_scale_workers(cfg, comm.get_world_size())