首页 > 其他分享> 文章详细

Pytorch基础操作 —— 3.保存和加载Torch模型和参数

2021-10-07 12:57:57 阅读：446 来源： 互联网

标签：load TorchScript torch 模型 Torch Pytorch model save 加载

文章目录

基础操作 —— 模型的存储
- S/L 模型
- S/L 参数
TorchScript
- Pytorch模型转成TorchScript
- C/C++加载TorchScript

基础操作 —— 模型的存储

我们的网络训练完成后，如果表现还不错，通常会想到把模型保存起来。所谓的模型，指的是我们用什么的神经层构建的网络，而与模型一同进行保存的，还有与模型相关的参数解。

所以从这个概念出发，我们可以有两种保存和读取方法。

S/L 模型

这种会把模型和相关参数一起保存起来，这种做法能很方便我们在做其他工程时，把之前的网络集成到新的应用里。

# save the model 
torch.save(the_model, PATH)

# load the model
the_model = torch.load(PATH)

但是它会有个缺点，就是保存的文件体积特别大，因为它不光保存了模型，也保留了相关参数。所以有时候我们会想，如果我们已知了某种网络结构，是否可以只保留参数？

S/L 参数

这便是上面提到的第二种方法，就是仅存储模型的训练参数，这种方法有个前置条件，就是要求使用者已知网络模型。

# save the arguments
torch.save(the_model.state_dict(), PATH)

# load the arguments
the_model = TheModelClass(*args, **kwargs)
the_model.load_state_dict(torch.load(PATH))

在某些时候，我们也可以使用这个方法来调整网络的预训练模型，比如在做YOLO的时候，也经常会用到这种类似的方法，但这一部分我就不在这篇文章里展开了。

TorchScript

以上两种方法，适用于Python对Python的项目，如果是Python对C/CPP，Python对Java的项目，就不能使用上面提到的两种方法。

这时我们会使用第三种，把模型以TorchScript的形式进行保存，并加载。比方说我们训练了一个网络模型，它能识别文字信息，现在我们需要把相关任务以TorchScript的形式进行保存，这样就可以集成到比如C/C++程序，或者Java程序里。

Pytorch模型转成TorchScript

	# define a neural network module
	model = DefinedNeuralNetworkModule

    # converting to Torch Script via Annotation
    serialized_model = torch.jit.script(model)

    # save the torch script for C++
    serialized_model.save("LSTM_Classfication.pt")

C/C++加载TorchScript

    torch::jit::script::Module module;
    try {
        // Deserialize the ScriptModule from a file using torch::jit::load().
        module = torch::jit::load(argv[1]);
    }
    catch (const c10::Error& e) {
        std::cerr << "error loading the model." << std::endl;
        return -1;
    }

标签：load,TorchScript,torch,模型,Torch,Pytorch,model,save,加载
来源： https://blog.csdn.net/poisonchry/article/details/120628550

本站声明： 1. iCode9 技术分享网（下文简称本站）提供的所有内容，仅供技术学习、探讨和分享；
2. 关于本站的所有留言、评论、转载及引用，纯属内容发起人的个人观点，与本站观点和立场无关；
3. 关于本站的所有言论和文字，纯属内容发起人的个人观点，与本站观点和立场无关；
4. 本站文章均是网友提供，不完全保证技术分享内容的完整性、准确性、时效性、风险性和版权归属；如您发现该文章侵犯了您的权益，可联系我们第一时间进行删除；
5. 本站为非盈利性的个人网站，所有内容不会用来进行牟利，也不会利用任何形式的广告来间接获益，纯粹是为了广大技术爱好者提供技术内容和技术思想的分享性交流网站。

ICode9