保存
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")