def __init__(self, cfg, optimization_level):
logger = logging.getLogger("detectron2")
if not logger.isEnabledFor(logging.INFO):
setup_logger()
model = self.build_model(cfg)
optimizer = self.build_optimizer(cfg, model)
model, optimizer = amp.initialize(model,
optimizer,
opt_level=optimization_level)
data_loader = self.build_train_loader(cfg)
SimpleTrainer.__init__(self, 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, weights: Union[str, Dict[str, Any]]):
self.start_iter = 0
self.max_iter = cfg.SOLVER.MAX_ITER
self.cfg = cfg
# We do not make any super call here and implement `__init__` from
# `DefaultTrainer`: we need to initialize mixed precision model before
# wrapping to DDP, so we need to do it this way.
model = self.build_model(cfg)
optimizer = self.build_optimizer(cfg, model)
data_loader = self.build_train_loader(cfg)
scheduler = self.build_lr_scheduler(cfg, optimizer)
# Load pre-trained weights before wrapping to DDP because `ApexDDP` has
# some weird issue with `DetectionCheckpointer`.
# fmt: off
if isinstance(weights, str):
# weights are ``str`` means ImageNet init or resume training.
self.start_iter = (DetectionCheckpointer(
model, optimizer=optimizer,
scheduler=scheduler).resume_or_load(weights, resume=True).get(
"iteration", -1) + 1)
elif isinstance(weights, dict):
# weights are a state dict means our pretrain init.
DetectionCheckpointer(model)._load_model(weights)
# fmt: on
# Enable distributed training if we have multiple GPUs. Use Apex DDP for
# non-FPN backbones because its `delay_allreduce` functionality helps with
# gradient checkpointing.
if dist.get_world_size() > 1:
if global_cfg.get("GRADIENT_CHECKPOINT", False):
model = ApexDDP(model, delay_allreduce=True)
else:
model = nn.parallel.DistributedDataParallel(
model,
device_ids=[dist.get_rank()],
broadcast_buffers=False)
# Call `__init__` from grandparent class: `SimpleTrainer`.
SimpleTrainer.__init__(self, model, data_loader, optimizer)
self.scheduler = scheduler
self.checkpointer = DetectionCheckpointer(model,
cfg.OUTPUT_DIR,
optimizer=optimizer,
scheduler=self.scheduler)
self.register_hooks(self.build_hooks())