在用C#调用pytorch模型时,报错,网上也有人报错,但几乎没有给出好的解决方案
问题分析:
用torch.save(model.state_dict(), ckpt_path)方法保存模型,只保存了参数,模型并没有保存下来
网上有人建议用torch.save(model,ckpt_path)保存,我试了还是不行
解决方案
traced_script_module = torch.jit.script(model, (x, y))
traced_script_module.save(ckpt_path)
注意
- (x,y)是模型传的参数,我的测试中模型需要两个参数。如果是一个就直接写(model,x)即可.
- 参数的类型要求也很严格,Union[Tensor, Tuple[Tensor], Dict[str, Tensor]]。如果参数类型错误也会报错
- 在保存的时候,我在调试过程中发现,它会将你给的参数传入模型的forward中执行一遍。