In the process of training the deep learning model, it is best to save the trained deep learning model, or conveniently load the model trained by others to save training time and achieve efficient problem solving.
Why do you need model files
- There are many model parameters for deep learning, such as: Transformer model, Bert model, etc.
- The training data set is generally very large, for example: 1000G or more.
- If the computing power of the local computer or the computing power of the server in the laboratory is basically insufficient, it will take a long time to train the model. If a model is short, the training cannot be stopped for several days, or even several months, and there may be insufficient memory. At this time, it is very good if there is a similar trained model that can be used directly and then fine-tuned.
torch.save() saves the network structure and model parameters. There are two ways to save
- One is to save the structural information and model parameter information of the entire neural network, and the object of save is the entire network model;
- The second is to save only the training model parameters of the neural network, and the object of save is net.state_dict().
Suppose there is a trained model named net1
torch.save(net1, ‘7-net.pth’) # 保存整个神经网络的结构和模型参数
torch.save(net1, ‘7-net.pkl’) # 同上
torch.save(net1.state_dict(), ‘7-net_params.pth’) # 只保存神经网络的模型参数
torch.save(net1.state_dict(), ‘7-net_params.pkl’) # 同上
If the torch.save method is used to save the model parameters, the suffix of the saved file has no effect, and the result is the same. Many .pkl files are also saved with torch.save, which is exactly the same as the .pth file.
-
However, if the application scenario is not here, there is still a difference between the two formats of files. The .pkl file is a format for saving files in python. If you open it directly, it will display a bunch of serialized things, which are actually in binary form. Stored, if you want to read these files, you need to use 'rb' instead of 'r' mode.
-
The .pth file has different applications. When Python traverses the known library file directory, if it sees a .pth file, it will add the path recorded in the file to the sys.path setting, so .pth The library specified by the file can also be found by the Python runtime environment.
But in fact, no matter pkl file or pth file, they are all stored in binary form, there is no essential difference. If you use the pickle library to load pkl file or pth file, the effect is the same.
Model file suffix
.pt: This suffix is often used in official documents..pth: This suffix is generally used by everyone as a convention..pkl: This suffix is because Python has a serialization module pickle, and when using it to save the model, it usually starts a file with the suffix .pkl.