【Fast-ReID】部署torch保存模型及参数

保存

cfg.defrost()
cfg.MODEL.BACKBONE.PRETRAIN = False
model = DefaultTrainer.build_model(cfg)
Checkpointer(model).load(cfg.MODEL.WEIGHTS)  # load trained model
model.eval()
inputs = torch.randn(1,3,cfg.INPUT.SIZE_TEST[0],cfg.INPUT.SIZE_TEST[1]).to(model.device)
scriptModel = torch.jit.trace(model,inputs)
scriptModel.save("XXX.pt")
 

加载

torch.jit.load("XXX.pt")