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 __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 __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 main(args): cfg = setup(args) if args.eval_only: cfg.defrost() cfg.MODEL.BACKBONE.PRETRAIN = args.imageNet model = DefaultTrainer.build_model(cfg) Checkpointer(model).load(cfg.MODEL.WEIGHTS) # load trained model res = DefaultTrainer.test(cfg, model) return res trainer = DefaultTrainer(cfg) if args.finetune: C = Checkpointer(trainer.model) C.load(cfg.MODEL.WEIGHTS) # load trained model to funetune trainer.resume_or_load(resume=args.resume) return trainer.train()