def main(args):
cfg = setup(args)
if args.eval_only:
cfg.defrost()
cfg.MODEL.BACKBONE.PRETRAIN = False
model = ClasTrainer.build_model(cfg)
Checkpointer(model).load(cfg.MODEL.WEIGHTS) # load trained model
try:
output_dir = os.path.dirname(cfg.MODEL.WEIGHTS)
path = os.path.join(output_dir, "idx2class.json")
with PathManager.open(path, 'r') as f:
idx2class = json.load(f)
ClasTrainer.idx2class = idx2class
except:
logger = logging.getLogger("fastreid.fastclas")
logger.info(
f"Cannot find idx2class dict in {os.path.dirname(cfg.MODEL.WEIGHTS)}"
)
res = ClasTrainer.test(cfg, model)
return res
trainer = ClasTrainer(cfg)
trainer.resume_or_load(resume=args.resume)
return trainer.train()
def __init__(self, cfg):
"""
Args:
cfg (CfgNode):
"""
super().__init__(cfg)
model_t = self.build_model_teacher(self.cfg)
for param in model_t.parameters():
param.requires_grad = False
logger = logging.getLogger('fastreid.' + __name__)
# Load pre-trained teacher model
print('88888\n' * 10)
logger.info("Loading teacher model ...")
Checkpointer(model_t).load(cfg.MODEL.TEACHER_WEIGHTS)
# Load pre-trained student model
# logger.info("Loading student model ...")
# Checkpointer(self.model, self.data_loader.dataset).load(cfg.MODEL.STUDENT_WEIGHTS)
def set_bn_eval(m):
classname = m.__class__.__name__
if classname.find('BatchNorm') != -1:
m.eval()
model_t.apply(set_bn_eval)
self.model_t = model_t
def __init__(self,
fastreid_config_path: str,
model_path: Optional[str],
use_cuda=True):
cfg = get_cfg()
cfg.merge_from_file(fastreid_config_path)
self.net = build_model(cfg)
self.device = "cuda" if torch.cuda.is_available(
) and use_cuda else "cpu"
logger = logging.getLogger("root.tracker")
self.net.to(self.device)
self.net.eval()
if model_path is not None:
logger.info(f"Loading weights from {model_path}")
Checkpointer(self.net).load(model_path)
else:
logger.info("Loading Model Pretrained on ImageNet")
height, width = cfg.INPUT.SIZE_TEST
self.size = (width, height)
logger.info(f"Image size: {self.size}")
self.norm = transforms.Compose([
transforms.ToTensor(),
transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]),
])
def main(args):
cfg = setup(args)
logger = logging.getLogger('fastreid.' + __name__)
if args.eval_only:
cfg.defrost()
cfg.MODEL.BACKBONE.PRETRAIN = False
model = Trainer.build_model(cfg)
model = nn.DataParallel(model)
model = model.cuda()
Checkpointer(model, save_dir=cfg.OUTPUT_DIR).load(cfg.MODEL.WEIGHTS) # load trained model
if cfg.TEST.PRECISE_BN.ENABLED and hooks.get_bn_modules(model):
prebn_cfg = cfg.clone()
prebn_cfg.DATALOADER.NUM_WORKERS = 0 # save some memory and time for PreciseBN
prebn_cfg.DATASETS.NAMES = tuple([cfg.TEST.PRECISE_BN.DATASET]) # set dataset name for PreciseBN
logger.info("Prepare precise BN dataset")
hooks.PreciseBN(
# Run at the same freq as (but before) evaluation.
model,
# Build a new data loader to not affect training
Trainer.build_train_loader(prebn_cfg),
cfg.TEST.PRECISE_BN.NUM_ITER,
).update_stats()
res = Trainer.test(cfg, model)
return res
trainer = Trainer(cfg)
trainer.resume_or_load(resume=args.resume)
return trainer.train()
def __init__(self, cfg):
"""
Args:
cfg (CfgNode):
"""
logger = logging.getLogger("fastreid")
if not logger.isEnabledFor(
logging.INFO): # setup_logger is not called for fastreid
setup_logger()
# Assume these objects must be constructed in this order.
data_loader = self.build_train_loader(cfg)
cfg = self.auto_scale_hyperparams(cfg, data_loader)
model = self.build_model(cfg)
optimizer = self.build_optimizer(cfg, model)
# For training, wrap with DDP. But don't need this for inference.
if comm.get_world_size() > 1:
# ref to https://github.com/pytorch/pytorch/issues/22049 to set `find_unused_parameters=True`
# for part of the parameters is not updated.
model = DistributedDataParallel(model,
device_ids=[comm.get_local_rank()],
broadcast_buffers=False)
super().__init__(model, data_loader, optimizer, cfg.SOLVER.AMP_ENABLED)
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 = Checkpointer(
# Assume you want to save checkpoints together with logs/statistics
model,
cfg.OUTPUT_DIR,
save_to_disk=comm.is_main_process(),
optimizer=optimizer,
scheduler=self.scheduler,
)
self.start_iter = 0
if cfg.SOLVER.SWA.ENABLED:
self.max_iter = cfg.SOLVER.MAX_ITER + cfg.SOLVER.SWA.ITER
else:
self.max_iter = cfg.SOLVER.MAX_ITER
self.cfg = cfg
self.register_hooks(self.build_hooks())
def __init__(self, cfg_):
self.cfg = cfg_.clone() # cfg can be modified by model
self.cfg.defrost()
self.cfg.MODEL.BACKBONE.PRETRAIN = False
self.model = build_model(self.cfg)
self.model.eval()
Checkpointer(self.model).load(cfg_.MODEL.WEIGHTS)
def main(args):
cfg = setup(args)
if args.eval_only:
cfg.defrost()
cfg.MODEL.BACKBONE.PRETRAIN = False
model = DefaultTrainer.build_model(cfg)
Checkpointer(model).load(cfg.MODEL.WEIGHTS) # load trained model
res = DefaultTrainer.test(cfg, model)
return res
trainer = DefaultTrainer(cfg)
# load trained model to funetune
if args.finetune: Checkpointer(trainer.model).load(cfg.MODEL.WEIGHTS)
trainer.resume_or_load(resume=args.resume)
return trainer.train()
def __init__(self, cfg):
"""
Args:
cfg (CfgNode):
"""
self.cfg = cfg
logger = logging.getLogger(__name__)
if not logger.isEnabledFor(
logging.INFO): # setup_logger is not called for fastreid
setup_logger()
# Assume these objects must be constructed in this order.
model = self.build_model(cfg)
optimizer = self.build_optimizer(cfg, model)
logger.info('Prepare training set')
data_loader = self.build_train_loader(cfg)
# For training, wrap with DP. But don't need this for inference.
model = DataParallel(model)
if cfg.MODEL.BACKBONE.NORM == "syncBN":
# Monkey-patching with syncBN
patch_replication_callback(model)
model = model.cuda()
super().__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 = Checkpointer(
# Assume you want to save checkpoints together with logs/statistics
model,
self.data_loader.loader.dataset,
cfg.OUTPUT_DIR,
optimizer=optimizer,
scheduler=self.scheduler,
)
self.start_iter = 0
if cfg.SOLVER.SWA.ENABLED:
self.max_iter = cfg.SOLVER.MAX_ITER + cfg.SOLVER.SWA.ITER
else:
self.max_iter = cfg.SOLVER.MAX_ITER
self.cfg = cfg
self.register_hooks(self.build_hooks())
def main(args):
cfg = setup(args)
logger = logging.getLogger("fastreid.trainer")
cfg.defrost()
cfg.MODEL.BACKBONE.PRETRAIN = False
model = H_Trainer.build_model(cfg)
Checkpointer(model).load(cfg.MODEL.WEIGHTS) # load trained model
res = H_Trainer.test(cfg, model)
print_csv_format(res)
def main(args):
cfg = setup(args)
if args.eval_only:
model = DefaultTrainer.build_model(cfg)
Checkpointer(model, save_dir=cfg.OUTPUT_DIR).load(cfg.MODEL.WEIGHTS)
res = DefaultTrainer.test(cfg, model)
return res
trainer = DefaultTrainer(cfg)
trainer.resume_or_load(resume=args.resume)
return trainer.train()
def __init__(self, cfg):
self.cfg = cfg.clone()
if cfg.MODEL.WEIGHTS.endswith('.pt'):
self.model = torch.jit.load(cfg.MODEL.WEIGHTS)
else:
self.model = build_model(cfg)
# load pre-trained model
Checkpointer(self.model).load(cfg.MODEL.WEIGHTS)
self.model.eval()
# self.model = nn.DataParallel(self.model)
self.model.cuda()
num_channels = len(cfg.MODEL.PIXEL_MEAN)
self.mean = torch.tensor(cfg.MODEL.PIXEL_MEAN).view(1, num_channels, 1, 1)
self.std = torch.tensor(cfg.MODEL.PIXEL_STD).view(1, num_channels, 1, 1)
def main(args):
cfg = setup(args)
if args.eval_only:
cfg.defrost()
cfg.MODEL.BACKBONE.PRETRAIN = False
model = Trainer.build_model(cfg)
model = nn.DataParallel(model)
model = model.cuda()
Checkpointer(model, save_dir=cfg.OUTPUT_DIR).load(cfg.MODEL.WEIGHTS)
res = Trainer.test(cfg, model)
return res
trainer = Trainer(cfg)
trainer.resume_or_load(resume=args.resume)
return trainer.train()
def main(args):
cfg = setup(args)
if args.eval_only:
cfg.defrost()
cfg.MODEL.BACKBONE.PRETRAIN = False
cfg.TEST.USE_ONLY_CLO = args.cconly
model = DefaultTrainer.build_model(cfg)
Checkpointer(model).load(cfg.MODEL.WEIGHTS) # load trained model
res = DefaultTrainer.test(cfg, model)
return res
trainer = DefaultTrainer(cfg)
trainer.resume_or_load(resume=args.resume)
return trainer.train()
def main(args):
cfg = setup(args)
if args.eval_only:
model = DefaultTrainer.build_model(cfg)
Checkpointer(model, save_dir=cfg.OUTPUT_DIR).resume_or_load(
cfg.MODEL.WEIGHTS, resume=args.resume)
res = DefaultTrainer.test(cfg, model)
return res
trainer = DefaultTrainer(cfg)
# moco pretrain
# import torch
# state_dict = torch.load('logs/model_0109999.pth')['model_ema']
# ret = trainer.model.module.load_state_dict(state_dict, strict=False)
#
trainer.resume_or_load(resume=args.resume)
return trainer.train()
def __init__(self, model_config, model_path, use_cuda=True):
cfg = get_cfg()
cfg.merge_from_file(model_config)
cfg.MODEL.BACKBONE.PRETRAIN = False
self.net = DefaultTrainer.build_model(cfg)
self.device = "cuda" if torch.cuda.is_available(
) and use_cuda else "cpu"
Checkpointer(self.net).load(model_path)
logger = logging.getLogger("root.tracker")
logger.info("Loading weights from {}... Done!".format(model_path))
self.net.to(self.device)
self.net.eval()
height, width = cfg.INPUT.SIZE_TEST
self.size = (width, height)
self.norm = transforms.Compose([
transforms.ToTensor(),
transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]),
])
def main(args):
cfg = setup(args)
model = build_model(cfg)
logger.info("Model:\n{}".format(model))
if args.eval_only:
cfg.defrost()
cfg.MODEL.BACKBONE.PRETRAIN = False
Checkpointer(model).load(cfg.MODEL.WEIGHTS) # load trained model
return do_test(cfg, model)
distributed = comm.get_world_size() > 1
if distributed:
model = DistributedDataParallel(model, delay_allreduce=True)
do_train(cfg, model, resume=args.resume)
return do_test(cfg, model)
def __init__(self, cfg):
super(Distiller, self).__init__(cfg)
# Get teacher model config
cfg_t = get_cfg()
cfg_t.merge_from_file(cfg.KD.MODEL_CONFIG)
model_t = build_model(cfg_t)
logger.info("Teacher model:\n{}".format(model_t))
# No gradients for teacher model
for param in model_t.parameters():
param.requires_grad_(False)
logger.info("Loading teacher model weights ...")
Checkpointer(model_t).load(cfg.KD.MODEL_WEIGHTS)
# Not register teacher model as `nn.Module`, this is
# make sure teacher model weights not saved
self.model_t = [model_t.backbone, model_t.heads]
def __init__(self, config_file):
cfg = get_cfg()
cfg.merge_from_file(config_file)
cfg.defrost()
cfg.MODEL.WEIGHTS = 'projects/bjzProject/logs/bjz/arcface_adam/model_final.pth'
model = build_model(cfg)
Checkpointer(model).resume_or_load(cfg.MODEL.WEIGHTS)
model.cuda()
model.eval()
self.model = model
# self.model = torch.jit.load("reid_model.pt")
# self.model.eval()
# self.model.cuda()
example = torch.rand(1, 3, 256, 128)
example = example.cuda()
traced_script_module = torch.jit.trace_module(model,
{'inference': example})
traced_script_module.save("reid_feat_extractor.pt")
def _do_eval(self):
results = self._func()
if results:
assert isinstance(
results, dict
), "Eval function must return a dict. Got {} instead.".format(results)
flattened_results = flatten_results_dict(results)
for k, v in flattened_results.items():
try:
v = float(v)
except Exception:
raise ValueError(
"[EvalHook] eval_function should return a nested dict of float. "
"Got '{}: {}' instead.".format(k, v)
)
# Remove extra memory cache of main process due to evaluation
torch.cuda.empty_cache()
self.step += 1
# Here we save a checkpoint. It is automatically registered with
# Ray Tune and will potentially be passed as the `checkpoint_dir`
# parameter in future iterations.
with tune.checkpoint_dir(step=self.step) as checkpoint_dir:
additional_state = {"iteration": int(self.trainer.iter)}
Checkpointer(
# Assume you want to save checkpoints together with logs/statistics
self.trainer.model,
checkpoint_dir,
save_to_disk=True,
optimizer=self.trainer.optimizer,
scheduler=self.trainer.scheduler,
).save(name="checkpoint", **additional_state)
metrics = dict(r1=results['Rank-1'], map=results['mAP'], score=(results['Rank-1'] + results['mAP']) / 2)
tune.report(**metrics)
def main(args):
cfg = setup(args)
if args.eval_only:
cfg.defrost()
cfg.MODEL.BACKBONE.PRETRAIN = False
model = Trainer.build_model(cfg)
Checkpointer(model, save_dir=cfg.OUTPUT_DIR).load(
cfg.MODEL.WEIGHTS) # load trained model
if args.commit:
res = Committer.test(cfg, model)
else:
res = Trainer.test(cfg, model)
return res
trainer = Trainer(cfg)
trainer.resume_or_load(resume=args.resume)
return trainer.train()
def __init__(self, cfg):
super().__init__(cfg)
# Get teacher model config
model_ts = []
for i in range(len(cfg.KD.MODEL_CONFIG)):
cfg_t = get_cfg()
cfg_t.merge_from_file(cfg.KD.MODEL_CONFIG[i])
model_t = build_model(cfg_t)
# No gradients for teacher model
for param in model_t.parameters():
param.requires_grad_(False)
logger.info("Loading teacher model weights ...")
Checkpointer(model_t).load(cfg.KD.MODEL_WEIGHTS[i])
model_ts.append(model_t)
# Not register teacher model as `nn.Module`, this is
# make sure teacher model weights not saved
self.model_ts = model_ts
class DefaultTrainer(SimpleTrainer):
"""
A trainer with default training logic. Compared to `SimpleTrainer`, it
contains the following logic in addition:
1. Create model, optimizer, scheduler, dataloader from the given config.
2. Load a checkpoint or `cfg.MODEL.WEIGHTS`, if exists.
3. Register a few common hooks.
It is created to simplify the **standard model training workflow** and reduce code boilerplate
for users who only need the standard training workflow, with standard features.
It means this class makes *many assumptions* about your training logic that
may easily become invalid in a new research. In fact, any assumptions beyond those made in the
:class:`SimpleTrainer` are too much for research.
The code of this class has been annotated about restrictive assumptions it mades.
When they do not work for you, you're encouraged to:
1. Overwrite methods of this class, OR:
2. Use :class:`SimpleTrainer`, which only does minimal SGD training and
nothing else. You can then add your own hooks if needed. OR:
3. Write your own training loop similar to `tools/plain_train_net.py`.
Also note that the behavior of this class, like other functions/classes in
this file, is not stable, since it is meant to represent the "common default behavior".
It is only guaranteed to work well with the standard models and training workflow in detectron2.
To obtain more stable behavior, write your own training logic with other public APIs.
Attributes:
scheduler:
checkpointer (DetectionCheckpointer):
cfg (CfgNode):
Examples:
.. code-block:: python
trainer = DefaultTrainer(cfg)
trainer.resume_or_load() # load last checkpoint or MODEL.WEIGHTS
trainer.train()
"""
def __init__(self, cfg):
"""
Args:
cfg (CfgNode):
"""
self.cfg = cfg
logger = logging.getLogger(__name__)
if not logger.isEnabledFor(
logging.INFO): # setup_logger is not called for fastreid
setup_logger()
# Assume these objects must be constructed in this order.
model = self.build_model(cfg)
optimizer = self.build_optimizer(cfg, model)
logger.info('Prepare training set')
data_loader = self.build_train_loader(cfg)
# For training, wrap with DP. But don't need this for inference.
model = DataParallel(model)
if cfg.MODEL.BACKBONE.NORM == "syncBN":
# Monkey-patching with syncBN
patch_replication_callback(model)
model = model.cuda()
super().__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 = Checkpointer(
# Assume you want to save checkpoints together with logs/statistics
model,
self.data_loader.loader.dataset,
cfg.OUTPUT_DIR,
optimizer=optimizer,
scheduler=self.scheduler,
)
self.start_iter = 0
if cfg.SOLVER.SWA.ENABLED:
self.max_iter = cfg.SOLVER.MAX_ITER + cfg.SOLVER.SWA.ITER
else:
self.max_iter = cfg.SOLVER.MAX_ITER
self.cfg = cfg
self.register_hooks(self.build_hooks())
def resume_or_load(self, resume=True):
"""
If `resume==True`, and last checkpoint exists, resume from it.
Otherwise, load a model specified by the config.
Args:
resume (bool): whether to do resume or not
"""
# The checkpoint stores the training iteration that just finished, thus we start
# at the next iteration (or iter zero if there's no checkpoint).
checkpoint = self.checkpointer.resume_or_load(self.cfg.MODEL.WEIGHTS,
resume=resume)
self.start_iter = checkpoint.get("iteration", -1) if resume else -1
# The checkpoint stores the training iteration that just finished, thus we start
# at the next iteration (or iter zero if there's no checkpoint).
self.start_iter += 1
# Prefetcher need to reset because it will preload a batch data, but we have updated
# dataset person identity dictionary.
self.data_loader.reset()
def build_hooks(self):
"""
Build a list of default hooks, including timing, evaluation,
checkpointing, lr scheduling, precise BN, writing events.
Returns:
list[HookBase]:
"""
logger = logging.getLogger(__name__)
cfg = self.cfg.clone()
cfg.defrost()
cfg.DATALOADER.NUM_WORKERS = 0 # save some memory and time for PreciseBN
cfg.DATASETS.NAMES = tuple([cfg.TEST.PRECISE_BN.DATASET
]) # set dataset name for PreciseBN
ret = [
hooks.IterationTimer(),
hooks.LRScheduler(self.optimizer, self.scheduler),
]
if cfg.SOLVER.SWA.ENABLED:
ret.append(
hooks.SWA(
cfg.SOLVER.MAX_ITER,
cfg.SOLVER.SWA.PERIOD,
cfg.SOLVER.SWA.LR_FACTOR,
cfg.SOLVER.SWA.ETA_MIN_LR,
cfg.SOLVER.SWA.LR_SCHED,
))
if cfg.TEST.PRECISE_BN.ENABLED and hooks.get_bn_modules(self.model):
logger.info("Prepare precise BN dataset")
ret.append(
hooks.PreciseBN(
# Run at the same freq as (but before) evaluation.
self.model,
# Build a new data loader to not affect training
self.build_train_loader(cfg),
cfg.TEST.PRECISE_BN.NUM_ITER,
))
if cfg.MODEL.OPEN_LAYERS != [''] and cfg.SOLVER.FREEZE_ITERS > 0:
open_layers = ",".join(cfg.MODEL.OPEN_LAYERS)
logger.info(
f'Open "{open_layers}" training for {cfg.SOLVER.FREEZE_ITERS:d} iters'
)
ret.append(
hooks.FreezeLayer(
self.model,
cfg.MODEL.OPEN_LAYERS,
cfg.SOLVER.FREEZE_ITERS,
))
# Do PreciseBN before checkpointer, because it updates the model and need to
# be saved by checkpointer.
# This is not always the best: if checkpointing has a different frequency,
# some checkpoints may have more precise statistics than others.
# if comm.is_main_process():
ret.append(
hooks.PeriodicCheckpointer(self.checkpointer,
cfg.SOLVER.CHECKPOINT_PERIOD))
def test_and_save_results():
self._last_eval_results = self.test(self.cfg, self.model)
return self._last_eval_results
# Do evaluation after checkpointer, because then if it fails,
# we can use the saved checkpoint to debug.
ret.append(hooks.EvalHook(cfg.TEST.EVAL_PERIOD, test_and_save_results))
# run writers in the end, so that evaluation metrics are written
ret.append(
hooks.PeriodicWriter(self.build_writers(), cfg.SOLVER.LOG_PERIOD))
return ret
def build_writers(self):
"""
Build a list of writers to be used. By default it contains
writers that write metrics to the screen,
a json file, and a tensorboard event file respectively.
If you'd like a different list of writers, you can overwrite it in
your trainer.
Returns:
list[EventWriter]: a list of :class:`EventWriter` objects.
It is now implemented by:
.. code-block:: python
return [
CommonMetricPrinter(self.max_iter),
JSONWriter(os.path.join(self.cfg.OUTPUT_DIR, "metrics.json")),
TensorboardXWriter(self.cfg.OUTPUT_DIR),
]
"""
# Assume the default print/log frequency.
return [
# It may not always print what you want to see, since it prints "common" metrics only.
CommonMetricPrinter(self.max_iter),
JSONWriter(os.path.join(self.cfg.OUTPUT_DIR, "metrics.json")),
TensorboardXWriter(self.cfg.OUTPUT_DIR),
]
def train(self):
"""
Run training.
Returns:
OrderedDict of results, if evaluation is enabled. Otherwise None.
"""
super().train(self.start_iter, self.max_iter)
# if hasattr(self, "_last_eval_results") and comm.is_main_process():
# verify_results(self.cfg, self._last_eval_results)
# return self._last_eval_results
@classmethod
def build_model(cls, cfg):
"""
Returns:
torch.nn.Module:
It now calls :func:`detectron2.modeling.build_model`.
Overwrite it if you'd like a different model.
"""
model = build_model(cfg)
logger = logging.getLogger(__name__)
logger.info("Model:\n{}".format(model))
return model
@classmethod
def build_optimizer(cls, cfg, model):
"""
Returns:
torch.optim.Optimizer:
It now calls :func:`detectron2.solver.build_optimizer`.
Overwrite it if you'd like a different optimizer.
"""
return build_optimizer(cfg, model)
@classmethod
def build_lr_scheduler(cls, cfg, optimizer):
"""
It now calls :func:`detectron2.solver.build_lr_scheduler`.
Overwrite it if you'd like a different scheduler.
"""
return build_lr_scheduler(cfg, optimizer)
@classmethod
def build_train_loader(cls, cfg):
"""
Returns:
iterable
It now calls :func:`fastreid.data.build_detection_train_loader`.
Overwrite it if you'd like a different data loader.
"""
return build_reid_train_loader(cfg)
@classmethod
def build_test_loader(cls, cfg, dataset_name):
"""
Returns:
iterable
It now calls :func:`detectron2.data.build_detection_test_loader`.
Overwrite it if you'd like a different data loader.
"""
return build_reid_test_loader(cfg, dataset_name)
@classmethod
def build_evaluator(cls, cfg, num_query, output_dir=None):
return ReidEvaluator(cfg, num_query, output_dir)
@classmethod
def test(cls, cfg, model, evaluators=None):
"""
Args:
cfg (CfgNode):
model (nn.Module):
evaluators (list[DatasetEvaluator] or None): if None, will call
:meth:`build_evaluator`. Otherwise, must have the same length as
`cfg.DATASETS.TEST`.
Returns:
dict: a dict of result metrics
"""
logger = logging.getLogger(__name__)
if isinstance(evaluators, DatasetEvaluator):
evaluators = [evaluators]
if evaluators is not None:
assert len(
cfg.DATASETS.TEST) == len(evaluators), "{} != {}".format(
len(cfg.DATASETS.TEST), len(evaluators))
results = OrderedDict()
for idx, dataset_name in enumerate(cfg.DATASETS.TESTS):
logger.info(f'prepare test set')
data_loader, num_query = cls.build_test_loader(cfg, dataset_name)
# When evaluators are passed in as arguments,
# implicitly assume that evaluators can be created before data_loader.
if evaluators is not None:
evaluator = evaluators[idx]
else:
try:
evaluator = cls.build_evaluator(cfg, num_query)
except NotImplementedError:
logger.warn(
"No evaluator found. Use `DefaultTrainer.test(evaluators=)`, "
"or implement its `build_evaluator` method.")
results[dataset_name] = {}
continue
results_i = inference_on_dataset(model, data_loader, evaluator)
results[dataset_name] = results_i
if comm.is_main_process():
assert isinstance(
results_i, dict
), "Evaluator must return a dict on the main process. Got {} instead.".format(
results_i)
logger.info("Evaluation results for {} in csv format:".format(
dataset_name))
print_csv_format(results_i)
if len(results) == 1:
results = list(results.values())[0]
return results
class DefaultTrainer(TrainerBase):
"""
A trainer with default training logic. Compared to `SimpleTrainer`, it
contains the following logic in addition:
1. Create model, optimizer, scheduler, dataloader from the given config.
2. Load a checkpoint or `cfg.MODEL.WEIGHTS`, if exists.
3. Register a few common hooks.
It is created to simplify the **standard model training workflow** and reduce code boilerplate
for users who only need the standard training workflow, with standard features.
It means this class makes *many assumptions* about your training logic that
may easily become invalid in a new research. In fact, any assumptions beyond those made in the
:class:`SimpleTrainer` are too much for research.
The code of this class has been annotated about restrictive assumptions it mades.
When they do not work for you, you're encouraged to:
1. Overwrite methods of this class, OR:
2. Use :class:`SimpleTrainer`, which only does minimal SGD training and
nothing else. You can then add your own hooks if needed. OR:
3. Write your own training loop similar to `tools/plain_train_net.py`.
Also note that the behavior of this class, like other functions/classes in
this file, is not stable, since it is meant to represent the "common default behavior".
It is only guaranteed to work well with the standard models and training workflow in fastreid.
To obtain more stable behavior, write your own training logic with other public APIs.
Attributes:
scheduler:
checkpointer:
cfg (CfgNode):
Examples:
.. code-block:: python
trainer = DefaultTrainer(cfg)
trainer.resume_or_load() # load last checkpoint or MODEL.WEIGHTS
trainer.train()
"""
def __init__(self, cfg):
"""
Args:
cfg (CfgNode):
"""
super().__init__()
logger = logging.getLogger("fastreid")
if not logger.isEnabledFor(
logging.INFO): # setup_logger is not called for fastreid
setup_logger()
# Assume these objects must be constructed in this order.
data_loader = self.build_train_loader(cfg)
cfg = self.auto_scale_hyperparams(cfg, data_loader.dataset.num_classes)
model = self.build_model(cfg)
optimizer = self.build_optimizer(cfg, model)
optimizer_ckpt = dict(optimizer=optimizer)
if cfg.SOLVER.FP16_ENABLED:
model, optimizer = amp.initialize(model, optimizer, opt_level="O1")
optimizer_ckpt.update(dict(amp=amp))
# For training, wrap with DDP. But don't need this for inference.
if comm.get_world_size() > 1:
# ref to https://github.com/pytorch/pytorch/issues/22049 to set `find_unused_parameters=True`
# for part of the parameters is not updated.
# model = DistributedDataParallel(
# model, device_ids=[comm.get_local_rank()], broadcast_buffers=False
# )
model = DistributedDataParallel(model, delay_allreduce=True)
self._trainer = (AMPTrainer if cfg.SOLVER.FP16_ENABLED else
SimpleTrainer)(model, data_loader, optimizer)
self.iters_per_epoch = len(
data_loader.dataset) // cfg.SOLVER.IMS_PER_BATCH
self.scheduler = self.build_lr_scheduler(cfg, optimizer,
self.iters_per_epoch)
# Assume no other objects need to be checkpointed.
# We can later make it checkpoint the stateful hooks
self.checkpointer = Checkpointer(
# Assume you want to save checkpoints together with logs/statistics
model,
cfg.OUTPUT_DIR,
save_to_disk=comm.is_main_process(),
**optimizer_ckpt,
**self.scheduler,
)
self.start_epoch = 0
self.max_epoch = cfg.SOLVER.MAX_EPOCH
self.max_iter = self.max_epoch * self.iters_per_epoch
self.warmup_iters = cfg.SOLVER.WARMUP_ITERS
self.delay_epochs = cfg.SOLVER.DELAY_EPOCHS
self.cfg = cfg
self.register_hooks(self.build_hooks())
def resume_or_load(self, resume=True):
"""
If `resume==True` and `cfg.OUTPUT_DIR` contains the last checkpoint (defined by
a `last_checkpoint` file), resume from the file. Resuming means loading all
available states (eg. optimizer and scheduler) and update iteration counter
from the checkpoint. ``cfg.MODEL.WEIGHTS`` will not be used.
Otherwise, this is considered as an independent training. The method will load model
weights from the file `cfg.MODEL.WEIGHTS` (but will not load other states) and start
from iteration 0.
Args:
resume (bool): whether to do resume or not
"""
# The checkpoint stores the training iteration that just finished, thus we start
# at the next iteration (or iter zero if there's no checkpoint).
checkpoint = self.checkpointer.resume_or_load(self.cfg.MODEL.WEIGHTS,
resume=resume)
if resume and self.checkpointer.has_checkpoint():
self.start_epoch = checkpoint.get("epoch", -1) + 1
# The checkpoint stores the training iteration that just finished, thus we start
# at the next iteration (or iter zero if there's no checkpoint).
def build_hooks(self):
"""
Build a list of default hooks, including timing, evaluation,
checkpointing, lr scheduling, precise BN, writing events.
Returns:
list[HookBase]:
"""
logger = logging.getLogger(__name__)
cfg = self.cfg.clone()
cfg.defrost()
cfg.DATALOADER.NUM_WORKERS = 0 # save some memory and time for PreciseBN
cfg.DATASETS.NAMES = tuple([cfg.TEST.PRECISE_BN.DATASET
]) # set dataset name for PreciseBN
ret = [
hooks.IterationTimer(),
hooks.LRScheduler(self.optimizer, self.scheduler),
]
# if cfg.SOLVER.SWA.ENABLED:
# ret.append(
# hooks.SWA(
# cfg.SOLVER.MAX_ITER,
# cfg.SOLVER.SWA.PERIOD,
# cfg.SOLVER.SWA.LR_FACTOR,
# cfg.SOLVER.SWA.ETA_MIN_LR,
# cfg.SOLVER.SWA.LR_SCHED,
# )
# )
if cfg.TEST.PRECISE_BN.ENABLED and hooks.get_bn_modules(self.model):
logger.info("Prepare precise BN dataset")
ret.append(
hooks.PreciseBN(
# Run at the same freq as (but before) evaluation.
self.model,
# Build a new data loader to not affect training
self.build_train_loader(cfg),
cfg.TEST.PRECISE_BN.NUM_ITER,
))
ret.append(
hooks.LayerFreeze(
self.model,
cfg.MODEL.FREEZE_LAYERS,
cfg.SOLVER.FREEZE_ITERS,
cfg.SOLVER.FREEZE_FC_ITERS,
))
# Do PreciseBN before checkpointer, because it updates the model and need to
# be saved by checkpointer.
# This is not always the best: if checkpointing has a different frequency,
# some checkpoints may have more precise statistics than others.
def test_and_save_results():
self._last_eval_results = self.test(self.cfg, self.model)
return self._last_eval_results
# Do evaluation before checkpointer, because then if it fails,
# we can use the saved checkpoint to debug.
ret.append(hooks.EvalHook(cfg.TEST.EVAL_PERIOD, test_and_save_results))
if comm.is_main_process():
ret.append(
hooks.PeriodicCheckpointer(self.checkpointer,
cfg.SOLVER.CHECKPOINT_PERIOD))
# run writers in the end, so that evaluation metrics are written
ret.append(hooks.PeriodicWriter(self.build_writers(), 200))
return ret
def build_writers(self):
"""
Build a list of writers to be used. By default it contains
writers that write metrics to the screen,
a json file, and a tensorboard event file respectively.
If you'd like a different list of writers, you can overwrite it in
your trainer.
Returns:
list[EventWriter]: a list of :class:`EventWriter` objects.
It is now implemented by:
.. code-block:: python
return [
CommonMetricPrinter(self.max_iter),
JSONWriter(os.path.join(self.cfg.OUTPUT_DIR, "metrics.json")),
TensorboardXWriter(self.cfg.OUTPUT_DIR),
]
"""
# Assume the default print/log frequency.
return [
# It may not always print what you want to see, since it prints "common" metrics only.
CommonMetricPrinter(self.max_iter),
JSONWriter(os.path.join(self.cfg.OUTPUT_DIR, "metrics.json")),
TensorboardXWriter(self.cfg.OUTPUT_DIR),
]
def train(self):
"""
Run training.
Returns:
OrderedDict of results, if evaluation is enabled. Otherwise None.
"""
super().train(self.start_epoch, self.max_epoch, self.iters_per_epoch)
if comm.is_main_process():
assert hasattr(self, "_last_eval_results"
), "No evaluation results obtained during training!"
return self._last_eval_results
def run_step(self):
self._trainer.iter = self.iter
self._trainer.run_step()
@classmethod
def build_model(cls, cfg):
"""
Returns:
torch.nn.Module:
It now calls :func:`fastreid.modeling.build_model`.
Overwrite it if you'd like a different model.
"""
model = build_model(cfg)
logger = logging.getLogger(__name__)
logger.info("Model:\n{}".format(model))
return model
@classmethod
def build_optimizer(cls, cfg, model):
"""
Returns:
torch.optim.Optimizer:
It now calls :func:`fastreid.solver.build_optimizer`.
Overwrite it if you'd like a different optimizer.
"""
return build_optimizer(cfg, model)
@classmethod
def build_lr_scheduler(cls, cfg, optimizer, iters_per_epoch):
"""
It now calls :func:`fastreid.solver.build_lr_scheduler`.
Overwrite it if you'd like a different scheduler.
"""
return build_lr_scheduler(cfg, optimizer, iters_per_epoch)
@classmethod
def build_train_loader(cls, cfg):
"""
Returns:
iterable
It now calls :func:`fastreid.data.build_reid_train_loader`.
Overwrite it if you'd like a different data loader.
"""
logger = logging.getLogger(__name__)
logger.info("Prepare training set")
return build_reid_train_loader(cfg)
@classmethod
def build_test_loader(cls, cfg, dataset_name):
"""
Returns:
iterable
It now calls :func:`fastreid.data.build_reid_test_loader`.
Overwrite it if you'd like a different data loader.
"""
return build_reid_test_loader(cfg, dataset_name)
@classmethod
def build_evaluator(cls, cfg, dataset_name, output_dir=None):
data_loader, num_query = cls.build_test_loader(cfg, dataset_name)
return data_loader, ReidEvaluator(cfg, num_query, output_dir)
@classmethod
def test(cls, cfg, model):
"""
Args:
cfg (CfgNode):
model (nn.Module):
Returns:
dict: a dict of result metrics
"""
logger = logging.getLogger(__name__)
results = OrderedDict()
for idx, dataset_name in enumerate(cfg.DATASETS.TESTS):
logger.info("Prepare testing set")
try:
data_loader, evaluator = cls.build_evaluator(cfg, dataset_name)
except NotImplementedError:
logger.warn(
"No evaluator found. implement its `build_evaluator` method."
)
results[dataset_name] = {}
continue
results_i = inference_on_dataset(model,
data_loader,
evaluator,
flip_test=cfg.TEST.FLIP_ENABLED)
results[dataset_name] = results_i
if comm.is_main_process():
assert isinstance(
results, dict
), "Evaluator must return a dict on the main process. Got {} instead.".format(
results)
print_csv_format(results)
if len(results) == 1: results = list(results.values())[0]
return results
@staticmethod
def auto_scale_hyperparams(cfg, num_classes):
r"""
This is used for auto-computation actual training iterations,
because some hyper-param, such as MAX_ITER, means training epochs rather than iters,
so we need to convert specific hyper-param to training iterations.
"""
cfg = cfg.clone()
frozen = cfg.is_frozen()
cfg.defrost()
# If you don't hard-code the number of classes, it will compute the number automatically
if cfg.MODEL.HEADS.NUM_CLASSES == 0:
output_dir = cfg.OUTPUT_DIR
cfg.MODEL.HEADS.NUM_CLASSES = num_classes
logger = logging.getLogger(__name__)
logger.info(
f"Auto-scaling the num_classes={cfg.MODEL.HEADS.NUM_CLASSES}")
# Update the saved config file to make the number of classes valid
if comm.is_main_process() and output_dir:
# Note: some of our scripts may expect the existence of
# config.yaml in output directory
path = os.path.join(output_dir, "config.yaml")
with PathManager.open(path, "w") as f:
f.write(cfg.dump())
if frozen: cfg.freeze()
return cfg
class DefaultTrainer(SimpleTrainer):
"""
A trainer with default training logic. Compared to `SimpleTrainer`, it
contains the following logic in addition:
1. Create model, optimizer, scheduler, dataloader from the given config.
2. Load a checkpoint or `cfg.MODEL.WEIGHTS`, if exists.
3. Register a few common hooks.
It is created to simplify the **standard model training workflow** and reduce code boilerplate
for users who only need the standard training workflow, with standard features.
It means this class makes *many assumptions* about your training logic that
may easily become invalid in a new research. In fact, any assumptions beyond those made in the
:class:`SimpleTrainer` are too much for research.
The code of this class has been annotated about restrictive assumptions it mades.
When they do not work for you, you're encouraged to:
1. Overwrite methods of this class, OR:
2. Use :class:`SimpleTrainer`, which only does minimal SGD training and
nothing else. You can then add your own hooks if needed. OR:
3. Write your own training loop similar to `tools/plain_train_net.py`.
Also note that the behavior of this class, like other functions/classes in
this file, is not stable, since it is meant to represent the "common default behavior".
It is only guaranteed to work well with the standard models and training workflow in fastreid.
To obtain more stable behavior, write your own training logic with other public APIs.
Attributes:
scheduler:
checkpointer:
cfg (CfgNode):
Examples:
.. code-block:: python
trainer = DefaultTrainer(cfg)
trainer.resume_or_load() # load last checkpoint or MODEL.WEIGHTS
trainer.train()
"""
def __init__(self, cfg):
"""
Args:
cfg (CfgNode):
"""
logger = logging.getLogger("fastreid")
if not logger.isEnabledFor(
logging.INFO): # setup_logger is not called for fastreid
setup_logger()
# Assume these objects must be constructed in this order.
data_loader = self.build_train_loader(cfg)
cfg = self.auto_scale_hyperparams(cfg, data_loader)
model = self.build_model(cfg)
optimizer = self.build_optimizer(cfg, model)
# For training, wrap with DDP. But don't need this for inference.
if comm.get_world_size() > 1:
# ref to https://github.com/pytorch/pytorch/issues/22049 to set `find_unused_parameters=True`
# for part of the parameters is not updated.
model = DistributedDataParallel(model,
device_ids=[comm.get_local_rank()],
broadcast_buffers=False)
super().__init__(model, data_loader, optimizer, cfg.SOLVER.AMP_ENABLED)
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 = Checkpointer(
# Assume you want to save checkpoints together with logs/statistics
model,
cfg.OUTPUT_DIR,
save_to_disk=comm.is_main_process(),
optimizer=optimizer,
scheduler=self.scheduler,
)
self.start_iter = 0
if cfg.SOLVER.SWA.ENABLED:
self.max_iter = cfg.SOLVER.MAX_ITER + cfg.SOLVER.SWA.ITER
else:
self.max_iter = cfg.SOLVER.MAX_ITER
self.cfg = cfg
self.register_hooks(self.build_hooks())
def resume_or_load(self, resume=True):
"""
If `resume==True`, and last checkpoint exists, resume from it.
Otherwise, load a model specified by the config.
Args:
resume (bool): whether to do resume or not
"""
# The checkpoint stores the training iteration that just finished, thus we start
# at the next iteration (or iter zero if there's no checkpoint).
checkpoint = self.checkpointer.resume_or_load(self.cfg.MODEL.WEIGHTS,
resume=resume)
# Reinitialize dataloader iter because when we update dataset person identity dict
# to resume training, DataLoader won't update this dictionary when using multiprocess
# because of the function scope.
if resume and self.checkpointer.has_checkpoint():
self.start_iter = checkpoint.get("iteration", -1) + 1
# The checkpoint stores the training iteration that just finished, thus we start
# at the next iteration (or iter zero if there's no checkpoint).
def build_hooks(self):
"""
Build a list of default hooks, including timing, evaluation,
checkpointing, lr scheduling, precise BN, writing events.
Returns:
list[HookBase]:
"""
logger = logging.getLogger(__name__)
cfg = self.cfg.clone()
cfg.defrost()
cfg.DATALOADER.NUM_WORKERS = 0 # save some memory and time for PreciseBN
cfg.DATASETS.NAMES = tuple([cfg.TEST.PRECISE_BN.DATASET
]) # set dataset name for PreciseBN
ret = [
hooks.IterationTimer(),
hooks.LRScheduler(self.optimizer, self.scheduler),
]
if cfg.SOLVER.SWA.ENABLED:
ret.append(
hooks.SWA(
cfg.SOLVER.MAX_ITER,
cfg.SOLVER.SWA.PERIOD,
cfg.SOLVER.SWA.LR_FACTOR,
cfg.SOLVER.SWA.ETA_MIN_LR,
cfg.SOLVER.SWA.LR_SCHED,
))
if cfg.TEST.PRECISE_BN.ENABLED and hooks.get_bn_modules(self.model):
logger.info("Prepare precise BN dataset")
ret.append(
hooks.PreciseBN(
# Run at the same freq as (but before) evaluation.
self.model,
# Build a new data loader to not affect training
self.build_train_loader(cfg),
cfg.TEST.PRECISE_BN.NUM_ITER,
))
if cfg.MODEL.FREEZE_LAYERS != [''] and cfg.SOLVER.FREEZE_ITERS > 0:
freeze_layers = ",".join(cfg.MODEL.FREEZE_LAYERS)
logger.info(
f'Freeze layer group "{freeze_layers}" training for {cfg.SOLVER.FREEZE_ITERS:d} iterations'
)
ret.append(
hooks.FreezeLayer(
self.model,
self.optimizer,
cfg.MODEL.FREEZE_LAYERS,
cfg.SOLVER.FREEZE_ITERS,
))
# Do PreciseBN before checkpointer, because it updates the model and need to
# be saved by checkpointer.
# This is not always the best: if checkpointing has a different frequency,
# some checkpoints may have more precise statistics than others.
if comm.is_main_process():
ret.append(
hooks.PeriodicCheckpointer(self.checkpointer,
cfg.SOLVER.CHECKPOINT_PERIOD))
def test_and_save_results():
self._last_eval_results = self.test(self.cfg, self.model)
return self._last_eval_results
# Do evaluation after checkpointer, because then if it fails,
# we can use the saved checkpoint to debug.
ret.append(hooks.EvalHook(cfg.TEST.EVAL_PERIOD, test_and_save_results))
if comm.is_main_process():
# run writers in the end, so that evaluation metrics are written
ret.append(hooks.PeriodicWriter(self.build_writers(), 200))
return ret
def build_writers(self):
"""
Build a list of writers to be used. By default it contains
writers that write metrics to the screen,
a json file, and a tensorboard event file respectively.
If you'd like a different list of writers, you can overwrite it in
your trainer.
Returns:
list[EventWriter]: a list of :class:`EventWriter` objects.
It is now implemented by:
.. code-block:: python
return [
CommonMetricPrinter(self.max_iter),
JSONWriter(os.path.join(self.cfg.OUTPUT_DIR, "metrics.json")),
TensorboardXWriter(self.cfg.OUTPUT_DIR),
]
"""
# Assume the default print/log frequency.
return [
# It may not always print what you want to see, since it prints "common" metrics only.
CommonMetricPrinter(self.max_iter),
JSONWriter(os.path.join(self.cfg.OUTPUT_DIR, "metrics.json")),
TensorboardXWriter(self.cfg.OUTPUT_DIR),
]
def train(self):
"""
Run training.
Returns:
OrderedDict of results, if evaluation is enabled. Otherwise None.
"""
super().train(self.start_iter, self.max_iter)
if comm.is_main_process():
assert hasattr(self, "_last_eval_results"
), "No evaluation results obtained during training!"
# verify_results(self.cfg, self._last_eval_results)
return self._last_eval_results
@classmethod
def build_model(cls, cfg):
"""
Returns:
torch.nn.Module:
It now calls :func:`fastreid.modeling.build_model`.
Overwrite it if you'd like a different model.
"""
model = build_model(cfg)
logger = logging.getLogger(__name__)
logger.info("Model:\n{}".format(model))
return model
@classmethod
def build_optimizer(cls, cfg, model):
"""
Returns:
torch.optim.Optimizer:
It now calls :func:`fastreid.solver.build_optimizer`.
Overwrite it if you'd like a different optimizer.
"""
return build_optimizer(cfg, model)
@classmethod
def build_lr_scheduler(cls, cfg, optimizer):
"""
It now calls :func:`fastreid.solver.build_lr_scheduler`.
Overwrite it if you'd like a different scheduler.
"""
return build_lr_scheduler(cfg, optimizer)
@classmethod
def build_train_loader(cls, cfg):
"""
Returns:
iterable
It now calls :func:`fastreid.data.build_detection_train_loader`.
Overwrite it if you'd like a different data loader.
"""
logger = logging.getLogger(__name__)
logger.info("Prepare training set")
return build_reid_train_loader(cfg)
@classmethod
def build_test_loader(cls, cfg, dataset_name):
"""
Returns:
iterable
It now calls :func:`fastreid.data.build_detection_test_loader`.
Overwrite it if you'd like a different data loader.
"""
return build_reid_test_loader(cfg, dataset_name)
@classmethod
def build_evaluator(cls, cfg, dataset_name, output_dir=None):
data_loader, num_query = cls.build_test_loader(cfg, dataset_name)
return data_loader, ReidEvaluator(cfg, num_query, output_dir)
@classmethod
def test(cls, cfg, model):
"""
Args:
cfg (CfgNode):
model (nn.Module):
Returns:
dict: a dict of result metrics
"""
logger = logging.getLogger(__name__)
results = OrderedDict()
for idx, dataset_name in enumerate(cfg.DATASETS.TESTS):
logger.info("Prepare testing set")
try:
data_loader, evaluator = cls.build_evaluator(cfg, dataset_name)
except NotImplementedError:
logger.warn(
"No evaluator found. implement its `build_evaluator` method."
)
results[dataset_name] = {}
continue
results_i = inference_on_dataset(model, data_loader, evaluator)
results[dataset_name] = results_i
if comm.is_main_process():
assert isinstance(
results, dict
), "Evaluator must return a dict on the main process. Got {} instead.".format(
results)
print_csv_format(results)
if len(results) == 1: results = list(results.values())[0]
return results
@staticmethod
def auto_scale_hyperparams(cfg, data_loader):
r"""
This is used for auto-computation actual training iterations,
because some hyper-param, such as MAX_ITER, means training epochs rather than iters,
so we need to convert specific hyper-param to training iterations.
"""
cfg = cfg.clone()
frozen = cfg.is_frozen()
cfg.defrost()
iters_per_epoch = len(data_loader.dataset) // cfg.SOLVER.IMS_PER_BATCH
cfg.MODEL.HEADS.NUM_CLASSES = data_loader.dataset.num_classes
cfg.SOLVER.MAX_ITER *= iters_per_epoch
cfg.SOLVER.WARMUP_ITERS *= iters_per_epoch
cfg.SOLVER.FREEZE_ITERS *= iters_per_epoch
cfg.SOLVER.DELAY_ITERS *= iters_per_epoch
for i in range(len(cfg.SOLVER.STEPS)):
cfg.SOLVER.STEPS[i] *= iters_per_epoch
cfg.SOLVER.SWA.ITER *= iters_per_epoch
cfg.SOLVER.SWA.PERIOD *= iters_per_epoch
ckpt_multiple = cfg.SOLVER.CHECKPOINT_PERIOD / cfg.TEST.EVAL_PERIOD
# Evaluation period must be divided by 200 for writing into tensorboard.
eval_num_mod = (200 - cfg.TEST.EVAL_PERIOD * iters_per_epoch) % 200
cfg.TEST.EVAL_PERIOD = cfg.TEST.EVAL_PERIOD * iters_per_epoch + eval_num_mod
# Change checkpoint saving period consistent with evaluation period.
cfg.SOLVER.CHECKPOINT_PERIOD = int(cfg.TEST.EVAL_PERIOD *
ckpt_multiple)
logger = logging.getLogger(__name__)
logger.info(
f"Auto-scaling the config to num_classes={cfg.MODEL.HEADS.NUM_CLASSES}, "
f"max_Iter={cfg.SOLVER.MAX_ITER}, wamrup_Iter={cfg.SOLVER.WARMUP_ITERS}, "
f"freeze_Iter={cfg.SOLVER.FREEZE_ITERS}, delay_Iter={cfg.SOLVER.DELAY_ITERS}, "
f"step_Iter={cfg.SOLVER.STEPS}, ckpt_Iter={cfg.SOLVER.CHECKPOINT_PERIOD}, "
f"eval_Iter={cfg.TEST.EVAL_PERIOD}.")
if frozen: cfg.freeze()
return cfg
default=[],
nargs=argparse.REMAINDER,
)
return parser
if __name__ == '__main__':
args = get_parser().parse_args()
cfg = setup_cfg(args)
cfg.defrost()
cfg.MODEL.BACKBONE.PRETRAIN = False
if cfg.MODEL.HEADS.POOL_LAYER == 'FastGlobalAvgPool':
cfg.MODEL.HEADS.POOL_LAYER = 'GlobalAvgPool'
cfg.MODEL.BACKBONE.WITH_NL = False
model = build_model(cfg)
Checkpointer(model).load(cfg.MODEL.WEIGHTS)
model.eval()
logger.info(model)
inputs = torch.randn(1, 3, cfg.INPUT.SIZE_TEST[0],
cfg.INPUT.SIZE_TEST[1]).to(
torch.device(cfg.MODEL.DEVICE))
PathManager.mkdirs(args.output)
pytorch_to_caffe.trans_net(model, inputs, args.name)
pytorch_to_caffe.save_prototxt(f"{args.output}/{args.name}.prototxt")
pytorch_to_caffe.save_caffemodel(f"{args.output}/{args.name}.caffemodel")
logger.info(f"Export caffe model in {args.output} sucessfully!")
def __init__(self, cfg):
TrainerBase.__init__(self)
logger = logging.getLogger('fastreid.partial-fc.trainer')
if not logger.isEnabledFor(
logging.INFO): # setup_logger is not called for fastreid
setup_logger()
# Assume these objects must be constructed in this order.
data_loader = self.build_train_loader(cfg)
cfg = self.auto_scale_hyperparams(cfg, data_loader.dataset.num_classes)
model = self.build_model(cfg)
optimizer = self.build_optimizer(cfg, model)
if cfg.MODEL.HEADS.PFC.ENABLED:
# fmt: off
feat_dim = cfg.MODEL.BACKBONE.FEAT_DIM
embedding_dim = cfg.MODEL.HEADS.EMBEDDING_DIM
num_classes = cfg.MODEL.HEADS.NUM_CLASSES
sample_rate = cfg.MODEL.HEADS.PFC.SAMPLE_RATE
cls_type = cfg.MODEL.HEADS.CLS_LAYER
scale = cfg.MODEL.HEADS.SCALE
margin = cfg.MODEL.HEADS.MARGIN
# fmt: on
# Partial-FC module
embedding_size = embedding_dim if embedding_dim > 0 else feat_dim
self.pfc_module = PartialFC(embedding_size, num_classes,
sample_rate, cls_type, scale, margin)
self.pfc_optimizer = self.build_optimizer(cfg, self.pfc_module)
# For training, wrap with DDP. But don't need this for inference.
if comm.get_world_size() > 1:
# ref to https://github.com/pytorch/pytorch/issues/22049 to set `find_unused_parameters=True`
# for part of the parameters is not updated.
model = DistributedDataParallel(model,
device_ids=[comm.get_local_rank()],
broadcast_buffers=False,
find_unused_parameters=True)
self._trainer = PFCTrainer(model, data_loader, optimizer, self.pfc_module, self.pfc_optimizer) \
if cfg.MODEL.HEADS.PFC.ENABLED else SimpleTrainer(model, data_loader, optimizer)
self.iters_per_epoch = len(
data_loader.dataset) // cfg.SOLVER.IMS_PER_BATCH
self.scheduler = self.build_lr_scheduler(cfg, optimizer,
self.iters_per_epoch)
if cfg.MODEL.HEADS.PFC.ENABLED:
self.pfc_scheduler = self.build_lr_scheduler(
cfg, self.pfc_optimizer, self.iters_per_epoch)
self.checkpointer = Checkpointer(
# Assume you want to save checkpoints together with logs/statistics
model,
cfg.OUTPUT_DIR,
save_to_disk=comm.is_main_process(),
optimizer=optimizer,
**self.scheduler,
)
if cfg.MODEL.HEADS.PFC.ENABLED:
self.pfc_checkpointer = Checkpointer(
self.pfc_module,
cfg.OUTPUT_DIR,
optimizer=self.pfc_optimizer,
**self.pfc_scheduler,
)
self.start_epoch = 0
self.max_epoch = cfg.SOLVER.MAX_EPOCH
self.max_iter = self.max_epoch * self.iters_per_epoch
self.warmup_iters = cfg.SOLVER.WARMUP_ITERS
self.delay_epochs = cfg.SOLVER.DELAY_EPOCHS
self.cfg = cfg
self.register_hooks(self.build_hooks())
def __init__(self, cfg):
super().__init__()
logger = logging.getLogger("fastreid")
if not logger.isEnabledFor(
logging.INFO): # if setup_logger is not called for fastreid
setup_logger()
logger.info("==> Load target-domain dataset")
self.tgt = tgt = self.load_dataset(cfg.DATASETS.TGT)
self.tgt_nums = len(tgt.train)
cfg = self.auto_scale_hyperparams(cfg, self.tgt_nums)
# Create model
self.model = self.build_model(cfg,
load_model=cfg.MODEL.PRETRAIN,
show_model=True,
use_dsbn=False)
# Optimizer
self.optimizer, self.param_wrapper = self.build_optimizer(
cfg, self.model)
# For training, wrap with DDP. But don't need this for inference.
if comm.get_world_size() > 1:
# ref to https://github.com/pytorch/pytorch/issues/22049 to set `find_unused_parameters=True`
# for part of the parameters is not updated.
self.model = DistributedDataParallel(
self.model,
device_ids=[comm.get_local_rank()],
broadcast_buffers=False,
find_unused_parameters=True)
# Learning rate scheduler
self.iters_per_epoch = cfg.SOLVER.ITERS
self.scheduler = self.build_lr_scheduler(cfg, self.optimizer,
self.iters_per_epoch)
# Assume no other objects need to be checkpointed.
# We can later make it checkpoint the stateful hooks
self.checkpointer = Checkpointer(
# Assume you want to save checkpoints together with logs/statistics
self.model,
cfg.OUTPUT_DIR,
save_to_disk=comm.is_main_process(),
optimizer=self.optimizer,
**self.scheduler,
)
self.start_epoch = 0
self.max_epoch = cfg.SOLVER.MAX_EPOCH
self.max_iter = self.max_epoch * self.iters_per_epoch
self.warmup_iters = cfg.SOLVER.WARMUP_ITERS
self.delay_epochs = cfg.SOLVER.DELAY_EPOCHS
self.cfg = cfg
self.register_hooks(self.build_hooks())
if cfg.SOLVER.AMP.ENABLED:
unsupported = "AMPTrainer does not support single-process multi-device training!"
if isinstance(self.model, DistributedDataParallel):
assert not (self.model.device_ids
and len(self.model.device_ids) > 1), unsupported
from torch.cuda.amp.grad_scaler import GradScaler
self.grad_scaler = GradScaler()
else:
self.grad_scaler = None
def __init__(self, cfg):
"""
Args:
cfg (CfgNode):
"""
super().__init__()
logger = logging.getLogger("fastreid")
if not logger.isEnabledFor(
logging.INFO): # setup_logger is not called for fastreid
setup_logger()
# Assume these objects must be constructed in this order.
data_loader = self.build_train_loader(cfg)
cfg = self.auto_scale_hyperparams(cfg, data_loader.dataset.num_classes)
self.model = self.build_model(cfg)
self.optimizer, self.param_wrapper = self.build_optimizer(
cfg, self.model)
# For training, wrap with DDP. But don't need this for inference.
if comm.get_world_size() > 1:
# ref to https://github.com/pytorch/pytorch/issues/22049 to set `find_unused_parameters=True`
# for part of the parameters is not updated.
self.model = DistributedDataParallel(
self.model,
device_ids=[comm.get_local_rank()],
broadcast_buffers=False,
)
self._data_loader_iter = iter(data_loader)
self.iters_per_epoch = len(
data_loader.dataset) // cfg.SOLVER.IMS_PER_BATCH
self.scheduler = self.build_lr_scheduler(cfg, self.optimizer,
self.iters_per_epoch)
# Assume no other objects need to be checkpointed.
# We can later make it checkpoint the stateful hooks
self.checkpointer = Checkpointer(
# Assume you want to save checkpoints together with logs/statistics
self.model,
cfg.OUTPUT_DIR,
save_to_disk=comm.is_main_process(),
optimizer=self.optimizer,
**self.scheduler,
)
self.start_epoch = 0
self.max_epoch = cfg.SOLVER.MAX_EPOCH
self.max_iter = self.max_epoch * self.iters_per_epoch
self.warmup_iters = cfg.SOLVER.WARMUP_ITERS
self.delay_epochs = cfg.SOLVER.DELAY_EPOCHS
self.cfg = cfg
self.register_hooks(self.build_hooks())
if cfg.SOLVER.AMP.ENABLED:
unsupported = f"[{self.__class__.__name__}] does not support single-process multi-device training!"
if isinstance(self.model, DistributedDataParallel):
assert not (self.model.device_ids
and len(self.model.device_ids) > 1), unsupported
from torch.cuda.amp.grad_scaler import GradScaler
self.grad_scaler = GradScaler()
else:
self.grad_scaler = None
def do_train(cfg, model, resume=False):
data_loader = build_reid_train_loader(cfg)
model.train()
optimizer = build_optimizer(cfg, model)
iters_per_epoch = len(data_loader.dataset) // cfg.SOLVER.IMS_PER_BATCH
scheduler = build_lr_scheduler(cfg, optimizer, iters_per_epoch)
checkpointer = Checkpointer(model,
cfg.OUTPUT_DIR,
save_to_disk=comm.is_main_process(),
optimizer=optimizer**scheduler)
start_epoch = (checkpointer.resume_or_load(
cfg.MODEL.WEIGHTS, resume=resume).get("epoch", -1) + 1)
iteration = start_iter = start_epoch * iters_per_epoch
max_epoch = cfg.SOLVER.MAX_EPOCH
max_iter = max_epoch * iters_per_epoch
warmup_iters = cfg.SOLVER.WARMUP_ITERS
delay_epochs = cfg.SOLVER.DELAY_EPOCHS
periodic_checkpointer = PeriodicCheckpointer(checkpointer,
cfg.SOLVER.CHECKPOINT_PERIOD,
max_epoch)
writers = ([
CommonMetricPrinter(max_iter),
JSONWriter(os.path.join(cfg.OUTPUT_DIR, "metrics.json")),
TensorboardXWriter(cfg.OUTPUT_DIR)
] if comm.is_main_process() else [])
# compared to "train_net.py", we do not support some hooks, such as
# accurate timing, FP16 training and precise BN here,
# because they are not trivial to implement in a small training loop
logger.info("Start training from epoch {}".format(start_epoch))
with EventStorage(start_iter) as storage:
for epoch in range(start_epoch, max_epoch):
storage.epoch = epoch
for data, _ in zip(data_loader, range(iters_per_epoch)):
storage.iter = iteration
loss_dict = model(data)
losses = sum(loss_dict.values())
assert torch.isfinite(losses).all(), loss_dict
loss_dict_reduced = {
k: v.item()
for k, v in comm.reduce_dict(loss_dict).items()
}
losses_reduced = sum(loss
for loss in loss_dict_reduced.values())
if comm.is_main_process():
storage.put_scalars(total_loss=losses_reduced,
**loss_dict_reduced)
optimizer.zero_grad()
losses.backward()
optimizer.step()
storage.put_scalar("lr",
optimizer.param_groups[0]["lr"],
smoothing_hint=False)
if iteration - start_iter > 5 and (
(iteration + 1) % 200 == 0 or iteration == max_iter - 1):
for writer in writers:
writer.write()
iteration += 1
if iteration <= warmup_iters:
scheduler["warmup_sched"].step()
# Write metrics after each epoch
for writer in writers:
writer.write()
if iteration > warmup_iters and (epoch + 1) >= delay_epochs:
scheduler["lr_sched"].step()
if (cfg.TEST.EVAL_PERIOD > 0
and (epoch + 1) % cfg.TEST.EVAL_PERIOD == 0
and epoch != max_iter - 1):
do_test(cfg, model)
# Compared to "train_net.py", the test results are not dumped to EventStorage
periodic_checkpointer.step(epoch)
def __init__(self, cfg):
super().__init__()
logger = logging.getLogger("fastreid")
if not logger.isEnabledFor(
logging.INFO): # if setup_logger is not called for fastreid
setup_logger()
# Create datasets
logger.info("==> Load source-domain dataset")
self.src = src = self.load_dataset(cfg.DATASETS.SRC)
self.src_pid_nums = src.get_num_pids(src.train)
logger.info("==> Load target-domain dataset")
self.tgt = tgt = self.load_dataset(cfg.DATASETS.TGT)
self.tgt_nums = len(tgt.train)
# Create model
self.model = self.build_model(cfg,
load_model=False,
show_model=False,
use_dsbn=True)
# Create hybrid memorys
self.hm = HybridMemory(num_features=cfg.MODEL.BACKBONE.FEAT_DIM,
num_samples=self.src_pid_nums + self.tgt_nums,
temp=cfg.MEMORY.TEMP,
momentum=cfg.MEMORY.MOMENTUM,
use_half=cfg.SOLVER.AMP.ENABLED).cuda()
# Initialize source-domain class centroids
logger.info(
"==> Initialize source-domain class centroids in the hybrid memory"
)
with inference_context(self.model), torch.no_grad():
src_train = self.build_dataset(cfg,
src.train,
is_train=False,
relabel=False,
with_mem_idx=False)
src_init_feat_loader = self.build_test_loader(cfg, src_train)
src_fname_feat_dict, _ = extract_features(self.model,
src_init_feat_loader)
src_feat_dict = collections.defaultdict(list)
for f, pid, _ in sorted(src.train):
src_feat_dict[pid].append(src_fname_feat_dict[f].unsqueeze(0))
src_centers = [
torch.cat(src_feat_dict[pid], 0).mean(0)
for pid in sorted(src_feat_dict.keys())
]
src_centers = torch.stack(src_centers, 0)
src_centers = F.normalize(src_centers, dim=1)
# Initialize target-domain instance features
logger.info(
"==> Initialize target-domain instance features in the hybrid memory"
)
with inference_context(self.model), torch.no_grad():
tgt_train = self.build_dataset(cfg,
tgt.train,
is_train=False,
relabel=False,
with_mem_idx=False)
tgt_init_feat_loader = self.build_test_loader(cfg, tgt_train)
tgt_fname_feat_dict, _ = extract_features(self.model,
tgt_init_feat_loader)
tgt_features = torch.cat([
tgt_fname_feat_dict[f].unsqueeze(0)
for f, _, _ in sorted(self.tgt.train)
], 0)
tgt_features = F.normalize(tgt_features, dim=1)
self.hm.features = torch.cat((src_centers, tgt_features), dim=0).cuda()
del (src_train, src_init_feat_loader, src_fname_feat_dict,
src_feat_dict, src_centers, tgt_train, tgt_init_feat_loader,
tgt_fname_feat_dict, tgt_features)
# Optimizer
self.optimizer, self.param_wrapper = self.build_optimizer(
cfg, self.model)
# For training, wrap with DDP. But don't need this for inference.
if comm.get_world_size() > 1:
# ref to https://github.com/pytorch/pytorch/issues/22049 to set `find_unused_parameters=True`
# for part of the parameters is not updated.
self.model = DistributedDataParallel(
self.model,
device_ids=[comm.get_local_rank()],
broadcast_buffers=False,
find_unused_parameters=True)
# Learning rate scheduler
self.iters_per_epoch = cfg.SOLVER.ITERS
self.scheduler = self.build_lr_scheduler(cfg, self.optimizer,
self.iters_per_epoch)
# Assume no other objects need to be checkpointed.
# We can later make it checkpoint the stateful hooks
self.checkpointer = Checkpointer(
# Assume you want to save checkpoints together with logs/statistics
self.model,
cfg.OUTPUT_DIR,
save_to_disk=comm.is_main_process(),
optimizer=self.optimizer,
**self.scheduler,
)
self.start_epoch = 0
self.max_epoch = cfg.SOLVER.MAX_EPOCH
self.max_iter = self.max_epoch * self.iters_per_epoch
self.warmup_iters = cfg.SOLVER.WARMUP_ITERS
self.delay_epochs = cfg.SOLVER.DELAY_EPOCHS
self.cfg = cfg
self.register_hooks(self.build_hooks())
if cfg.SOLVER.AMP.ENABLED:
unsupported = "AMPTrainer does not support single-process multi-device training!"
if isinstance(self.model, DistributedDataParallel):
assert not (self.model.device_ids
and len(self.model.device_ids) > 1), unsupported
from torch.cuda.amp.grad_scaler import GradScaler
self.grad_scaler = GradScaler()
else:
self.grad_scaler = None
