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]) cfg_t.defrost() cfg_t.MODEL.META_ARCHITECTURE = "Baseline" # Change syncBN to BN due to no DDP wrapper if cfg_t.MODEL.BACKBONE.NORM == "syncBN": cfg_t.MODEL.BACKBONE.NORM = "BN" if cfg_t.MODEL.HEADS.NORM == "syncBN": cfg_t.MODEL.HEADS.NORM = "BN" 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) self.ema_enabled = cfg.KD.EMA.ENABLED self.ema_momentum = cfg.KD.EMA.MOMENTUM if self.ema_enabled: cfg_self = cfg.clone() cfg_self.defrost() cfg_self.MODEL.META_ARCHITECTURE = "Baseline" if cfg_self.MODEL.BACKBONE.NORM == "syncBN": cfg_self.MODEL.BACKBONE.NORM = "BN" if cfg_self.MODEL.HEADS.NORM == "syncBN": cfg_self.MODEL.HEADS.NORM = "BN" model_self = build_model(cfg_self) # No gradients for self model for param in model_self.parameters(): param.requires_grad_(False) if cfg_self.MODEL.WEIGHTS != '': logger.info("Loading self distillation model weights ...") Checkpointer(model_self).load(cfg_self.MODEL.WEIGHTS) else: # Make sure the initial state is same for param_q, param_k in zip(self.parameters(), model_self.parameters()): param_k.data.copy_(param_q.data) model_ts.insert(0, model_self) # Not register teacher model as `nn.Module`, this is # make sure teacher model weights not saved self.model_ts = model_ts
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 __init__(self, cfg): self.cfg = cfg.clone() # cfg can be modified by model self.model = build_model(self.cfg) self.model.eval() # self.metadata = MetadataCatalog.get(cfg.DATASETS.TEST[0]) checkpointer = Checkpointer(self.model) checkpointer.load(cfg.MODEL.WEIGHTS)
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
def __init__(self, cfg, device='cpu'): self.cfg = cfg.clone() # cfg can be modified by model self.cfg.defrost() self.cfg.MODEL.BACKBONE.PRETRAIN = False self.device = device self.model = build_model(self.cfg) self.model.to(device) self.model.eval() checkpointer = Checkpointer(self.model) checkpointer.load(cfg.MODEL.WEIGHTS)
def build_model(cls, cfg, show_model=True): """ 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) if show_model: logger = logging.getLogger('fastreid') logger.info("Model:\n{}".format(model)) return model
def __init__(self, cfg): self.cfg = cfg.clone() # cfg can be modified by model model = build_model(self.cfg) self.model = DataParallel(model) self.model.cuda() self.model.eval() checkpointer = Checkpointer(self.model) checkpointer.load(cfg.MODEL.WEIGHTS) 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 build_model(cls, cfg, load_model=True, show_model=True, use_dsbn=False): cfg = cfg.clone() # cfg can be modified by model cfg.defrost() cfg.MODEL.DEVICE = "cpu" model = build_model(cfg) logger = logging.getLogger('fastreid') if load_model: pretrain_path = cfg.MODEL.PRETRAIN_PATH try: state_dict = torch.load( pretrain_path, map_location=torch.device("cpu"))['model'] for layer in cfg.MODEL.IGNORE_LAYERS: if layer in state_dict.keys(): del state_dict[layer] logger.info(f"Loading pretrained model from {pretrain_path}") except FileNotFoundError as e: logger.info( f"{pretrain_path} is not found! Please check this path.") raise e except KeyError as e: logger.info( "State dict keys error! Please check the state dict.") raise e incompatible = model.load_state_dict(state_dict, strict=False) if incompatible.missing_keys: logger.info( get_missing_parameters_message(incompatible.missing_keys)) if incompatible.unexpected_keys: logger.info( get_unexpected_parameters_message( incompatible.unexpected_keys)) if use_dsbn: logger.info("==> Convert BN to Domain Specific BN") convert_dsbn(model) if show_model: logger.info("Model:\n{}".format(model)) model.to(torch.device("cuda")) return 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, 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
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 build_model_teacher(cls, cfg) -> nn.Module: cfg_t = update_model_teacher_config(cfg) model_t = build_model(cfg_t) return model_t