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

刺骨的言语ヽ痛彻心扉 2022-08-28 09:45 271阅读 0赞

### 文章目录 ###

*  基础操作 —— 模型的存储
 *   *  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;
        }