Pytorch保存模型

前提

模型训练好后自然想要将里面所有层涉及的权重保存下来，这样子我们的模型就能部署在任意有pytorch环境下了。

Torch.save/load

先介绍一下纯py接口的保存方式。

class my_model(nn.Modules):
	def __init__(self):
		super(my_module,self).__init__()
		self.relu = nn.Relu()
	def forward(x, self):
		return self.relu(x)
		......
		......
model = my_model()
torch.save(model.state_dict(),"./model_name.pth")
torch.load(model,"./model_name.pth)

这样就会在当前目录保存一份.pth文件了(里面只保存了这个模型的所有权重即.parameters())。
下面的load就是在其他脚本中使用这个模型预训练好的权重
Torch.save 官网有详细介绍

Torch.jit

这个方式的保存更加高级（保存为TorchScript）可以与Torch c++接口通用。
这样带来的好处就是保存下来的模型为编译过后的运行时不需要python解释器，运行速度会更快。并且这种方式可以连带模型的定义一起保存，无需import model。

一般有两种保存方式

1. torch.jit