def build_model(self, 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, sample_weights=self.sample_weights)
logger = logging.getLogger("fastreid.attr_model")
logger.info("Model:\n{}".format(model))
return model
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)
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, 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 build_model(cls, cfg):
model = build_model(cfg)
return model