| 35 | metadata[newkey] = metadata.pop(key) |
| 36 | |
| 37 | class SGTCheckPointer(Checkpointer): |
| 38 | def __init__(self, model, save_dir="", *, save_to_disk=None, **checkpointables): |
| 39 | if isinstance(model, DistributedDataParallel): |
| 40 | model = model.module |
| 41 | # super().__init__(model.detector, save_dir, **kwargs) |
| 42 | # tracker = model.tracker |
| 43 | is_main_process = comm.is_main_process() |
| 44 | super().__init__( |
| 45 | model, |
| 46 | save_dir, |
| 47 | save_to_disk=is_main_process if save_to_disk is None else save_to_disk, |
| 48 | **checkpointables, |
| 49 | ) |
| 50 | self.whole_weight_flag = False |
| 51 | self.backbone_name = self.get_backbone_name() # default by dla |
| 52 | backbone_convert_fn_dict = { |
| 53 | 'resnet': self._convert_weight_name_resnet, |
| 54 | 'dla': self._convert_weight_name_dla, |
| 55 | 'hourglass': self._convert_weight_name_hourglass, |
| 56 | } |
| 57 | self.convert_weight_name_fn = backbone_convert_fn_dict[self.backbone_name] |
| 58 | self.dcnv2_flag = False |
| 59 | |
| 60 | def get_backbone_name(self): |
| 61 | backbone_name = 'dla' |
| 62 | backbone_class_name = type(self.model.detector.backbone).__name__.lower() |
| 63 | if 'dla' in backbone_class_name: |
| 64 | backbone_name = 'dla' |
| 65 | elif 'resnet' in backbone_class_name: |
| 66 | backbone_name = 'resnet' |
| 67 | elif 'hourglass' in backbone_class_name: |
| 68 | backbone_name = 'hourglass' |
| 69 | else: |
| 70 | raise NotImplementedError(f"Loading backbone {backbone_class_name} is not yet supported") |
| 71 | return backbone_name |
| 72 | |
| 73 | def resume_or_load(self, paths: Dict[str, str], *, resume: bool = True) -> Dict[str, Any]: |
| 74 | if paths['total'] != '': |
| 75 | self.whole_weight_flag = True |
| 76 | path = paths['total'] |
| 77 | else: |
| 78 | assert paths['detector'] != '', "Specify either detector weight or total model weight" |
| 79 | path = paths['detector'] |
| 80 | if resume and self.has_checkpoint(): |
| 81 | path = self.get_checkpoint_file() |
| 82 | return self.load(path) |
| 83 | else: |
| 84 | return self.load(path, checkpointables=[]) |
| 85 | |
| 86 | def load(self, path: str, checkpointables: Optional[List[str]] = None) -> Dict[str, Any]: |
| 87 | if not path: |
| 88 | # no checkpoint provided |
| 89 | self.logger.info("No checkpoint found. Initializing model from scratch") |
| 90 | return {} |
| 91 | self.logger.info("Loading checkpoint from {}".format(path)) |
| 92 | if not os.path.isfile(path): |
| 93 | path = self.path_manager.get_local_path(path) |
| 94 | assert os.path.isfile(path), "Checkpoint {} not found!".format(path) |
no outgoing calls
no test coverage detected