Shulou(Shulou.com)06/03 Report--

This article focuses on "what is the process of saving and loading the PyTorch deep learning model". Interested friends may wish to take a look. The method introduced in this paper is simple, fast and practical. Let the editor take you to learn "what is the process of saving and loading the PyTorch deep learning model"!

1. Saving and loading of model parameters

Torch.save (module.state_dict (), path): use the module.state_dict () function to get the parameters and buffers that have been trained at each layer, and then save the parameters and buffers to the file storage path specified by path (commonly used file formats are .pt, .pth, or .pkl).

Torch.nn.Module.load_state_dict (state_dict): loads parameters and buffers from state_dict into Module and its subclasses.

The torch.nn.Module.state_dict () function returns an OrderedDict type dictionary object in python, which maps each layer to its corresponding parameter and buffer, whose key value is the name of the parameter or buffer. Only those layers whose parameters can be trained are saved to OrderedDict, such as convolution layer, linear layer, and so on.

The dictionary class in Python accesses data in "key: value" mode, and OrderedDict is a subclass of it, which enables the sorting of elements in the dictionary object (OrderedDict sorts elements according to the order in which they are placed). Because of the sorting, two OrderedDict dictionary objects in different order are treated as two different objects.

Example:

Import torchimport torch.nn as nnclass Net (nn.Module): def _ init__ (self): super (Net, self). _ _ init__ () self.conv1 = nn.Conv2d (1,2,3) self.pool1 = nn.MaxPool2d (2,2) def forward (self X): X = self.conv1 (x) x = self.pool1 (x) return x # initialize the network net = Net () net.conv1.weight [0] .detach (). Fill_ (1) net.conv1.weight [1] .detach (). Fill_ (2) net.conv1.bias.data.detach (). Zero_ () # to get state_dictstate_dict = net.state_dict () # the default is to traverse key So param_tensor is actually the key value for param_tensor in state_dict: print (param_tensor,':\ nrecording state _ tensor [param _ tensor]) # Save the model parameter torch.save (state_dict, "net_params.pth") # get the model parameter net.load_state_dict (state_dict) output by loading state_dict:

II. Preservation and loading of complete models

Torch.save (module, path): save the entire trained network model module to the file storage path specified by path (commonly used file format is .pt or .pth).

Torch.load (path): loads the entire neural network model saved to path.

Example:

Import torchimport torch.nn as nnclass Net (nn.Module): def _ init__ (self): super (Net, self). _ _ init__ () self.conv1 = nn.Conv2d (1,2,3) self.pool1 = nn.MaxPool2d (2,2) def forward (self X): X = self.conv1 (x) x = self.pool1 (x) return x # initialize the network net = Net () net.conv1.weight [0] .detach (). Fill_ (1) net.conv1.weight [1] .detach (). Fill_ (2) net.conv1.bias.data.detach (). Zero_ () # saves the entire network torch.save (net "net.pth") # load network net = torch.load ("net.pth") here I believe that you have a deeper understanding of "what is the process of saving and loading PyTorch deep learning model". You might as well do it in practice. Here is the website, more related content can enter the relevant channels to inquire, follow us, continue to learn!

Welcome to subscribe "Shulou Technology Information " to get latest news, interesting things and hot topics in the IT industry, and controls the hottest and latest Internet news, technology news and IT industry trends.