def __init__(self, cfg: CfgNode): """ Initialize chart-based loss from configuration options Args: cfg (CfgNode): configuration options """ # fmt: off self.heatmap_size = cfg.MODEL.ROI_DENSEPOSE_HEAD.HEATMAP_SIZE self.w_points = cfg.MODEL.ROI_DENSEPOSE_HEAD.POINT_REGRESSION_WEIGHTS self.w_part = cfg.MODEL.ROI_DENSEPOSE_HEAD.PART_WEIGHTS self.use_part_focal_loss = cfg.MODEL.ROI_DENSEPOSE_HEAD.PART_FOCAL_LOSS self.w_segm = cfg.MODEL.ROI_DENSEPOSE_HEAD.INDEX_WEIGHTS self.w_body = cfg.MODEL.ROI_DENSEPOSE_HEAD.BODY_WEIGHTS self.n_segm_chan = cfg.MODEL.ROI_DENSEPOSE_HEAD.NUM_COARSE_SEGM_CHANNELS self.w_smooth = cfg.MODEL.ROI_DENSEPOSE_HEAD.SMOOTH_WEIGHTS self.w_tv = cfg.MODEL.ROI_DENSEPOSE_HEAD.TV_WEIGHTS # fmt: on self.segm_trained_by_masks = cfg.MODEL.ROI_DENSEPOSE_HEAD.COARSE_SEGM_TRAINED_BY_MASKS # <<<<<<< HEAD self.use_mean_uv = cfg.MODEL.ROI_DENSEPOSE_HEAD.MEAN_UV_LOSS if self.use_part_focal_loss: gamma = cfg.MODEL.ROI_DENSEPOSE_HEAD.PART_FOCAL_GAMMA self.focal_loss = FocalLoss(gamma=gamma, alpha=None, size_average=False) self.use_teacher_student = cfg.MODEL.TEACHER_STUDENT self.teacher_cfg = cfg.MODEL.TEACHER_CFG_FILE self.teacher_weights = cfg.MODEL.TEACHER_WEIGHTS self.teach_ins_wo_gt_dp = cfg.MODEL.TEACH_INS_WO_GT_DP self.w_part_teach = cfg.MODEL.TEACH_PART_WEIGHTS self.w_points_teach = cfg.MODEL.TEACH_POINT_REGRESSION_WEIGHTS if self.use_teacher_student: from densepose.engine import Trainer from densepose.modeling.densepose_checkpoint import DensePoseCheckpointer from densepose.config import get_cfg, add_densepose_config self.teacher_cfg = get_cfg() add_densepose_config(self.teacher_cfg) self.teacher_cfg.merge_from_file(cfg.MODEL.TEACHER_CFG_FILE) self.teacher_model = Trainer.build_model(self.teacher_cfg) # pdb.set_trace() DensePoseCheckpointer(self.teacher_model).load( cfg.MODEL.TEACHER_WEIGHTS) self.teacher_model.eval() self.use_aux_global_s = cfg.MODEL.CONDINST.AUX_SUPERVISION_GLOBAL_S self.use_aux_global_skeleton = cfg.MODEL.CONDINST.AUX_SUPERVISION_GLOBAL_SKELETON self.use_aux_body_semantics = cfg.MODEL.CONDINST.AUX_SUPERVISION_BODY_SEMANTICS self.w_aux_global_s = cfg.MODEL.CONDINST.AUX_SUPERVISION_GLOBAL_S_WEIGHTS self.w_aux_global_skeleton = cfg.MODEL.CONDINST.AUX_SUPERVISION_GLOBAL_SKELETON_WEIGHTS self.w_aux_body_semantics = cfg.MODEL.CONDINST.AUX_SUPERVISION_BODY_SEMANTICS self.pred_ins_body = cfg.MODEL.CONDINST.PREDICT_INSTANCE_BODY # def __call__( # self, proposals_with_gt: List[Instances], densepose_predictor_outputs: Any, images=None # ======= self.segm_loss = MaskOrSegmentationLoss(cfg)
def main(args): cfg = setup(args) # disable strict kwargs checking: allow one to specify path handle # hints through kwargs, like timeout in DP evaluation PathManager.set_strict_kwargs_checking(False) if args.eval_only: model = Trainer.build_model(cfg) DetectionCheckpointer(model, save_dir=cfg.OUTPUT_DIR).resume_or_load( cfg.MODEL.WEIGHTS, resume=args.resume ) res = Trainer.test(cfg, model) if cfg.TEST.AUG.ENABLED: res.update(Trainer.test_with_TTA(cfg, model)) if comm.is_main_process(): verify_results(cfg, res) return res trainer = Trainer(cfg) trainer.resume_or_load(resume=args.resume) if cfg.TEST.AUG.ENABLED: trainer.register_hooks( [hooks.EvalHook(0, lambda: trainer.test_with_TTA(cfg, trainer.model))] ) return trainer.train()