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

How is the saving and loading of PyTorch model? in view of this problem, this article introduces the corresponding analysis and solution in detail, hoping to help more partners who want to solve this problem to find a more simple and easy way.

Torch.save () and torch.load ():

Torch.save () and torch.load () are used together to save an object (any object, not necessarily an object in PyTorch) to a file, and to load an object from a file. When loading, you can indicate whether data is required to move between CPU and GPU.

Module.state_dict () and Module.load_state_dict ():

Module.state_dict () returns a dictionary that holds the entire state of Module as key-value pairs. Module.load_state_dict () can load parameters from a dictionary into the module and its descendants. If strict is True, then the loaded dictionary and the keywords returned by the state_dict () method of the module must match exactly .if strict is True Then the keys of state_dict must exactly match the keys returned by this module's state_dict () function. The return value is a named tuple: NamedTuple with missing_keys and unexpected_keys fields, which holds missing and unexpected keywords, respectively. If only part of the layer of your model is the same as that of the pre-training model, you can load only these same parameters, as long as you set the strict parameter to False to ignore those keys that do not match. # Mode 1 model_path = 'model_name.pth'# model_params_path =' params_name.pth'#-# torch.save (model, model_path) #-load-# model = torch.load (model_path) # Mode 2-# torch.save (model.state_dict ()) Model_params_path) # saved file name suffix is usually .pt or .pth #-load-# model=Model (). Cuda () # define model structure # model.load_state_dict (torch.load (model_params_path)) # load model parameters

Description:

# Save / load the entire model torch.save (model, PATH) model = torch.load (PATH) model.eval () this process of saving / loading the model uses the most intuitive syntax and uses less code. This uses Python's pickle to save all modules. The disadvantage of this approach is that when the model is saved, the serialized data is bound to a specific class and exact directory. This is because pickle does not save the model class itself, but saves the path to the class and uses it when loading. Therefore, when using or refactoring in other projects, there will be errors when loading the model. # Save / load state_dict (recommended) torch.save (model.state_dict (), PATH) model = TheModelClass (* args, * * kwargs) model.load_state_dict (torch.load (PATH)) model.eval ()

Select the parameters to save and set checkpoint:

#-Save-torch.save ({'epoch': epoch + 1 grammatical archetypes: args.arch,'state_dict': model.state_dict (),' optimizer_state_dict': optimizer.state_dict (), 'loss': loss,'best_prec1': best_prec1,} 'checkpoint_name.tar') #-load-- checkpoint = torch.load (' checkpoint_name.tar') # get the saved parameter start_epoch = checkpoint ['epoch'] best_prec1 = checkpoint [' best_prec1'] state_dict=checkpoint ['state_dict'] model=Model () # define the model structure model.load_state_dict (state_dict)

Save multiple models to the same file:

#-Save-torch.save ({'modelA_state_dict': modelA.state_dict (),' modelB_state_dict': modelB.state_dict (), 'optimizerA_state_dict': optimizerA.state_dict (),' optimizerB_state_dict': optimizerB.state_dict (),...}, PATH) #-modelA = TheModelAClass (* args, * * kwargs) modelB = TheModelAClass (* args * * kwargs) optimizerA = TheOptimizerAClass (* args, * * kwargs) optimizerB = TheOptimizerBClass (* args * * kwargs) checkpoint = torch.load (PATH) modelA.load_state_dict (checkpoint ['modelA_state_dict'] modelB.load_state_dict (checkpoint [' modelB_state_dict'] optimizerA.load_state_dict (checkpoint ['optimizerA_state_dict'] optimizerB.load_state_dict (checkpoint [' optimizerB_state_dict'] modelA.eval () modelB.eval () # ormodelA.train () modelB.train () # here it is When loading after saving the model, I sometimes # encounter the problem of CUDA out of memory, # the solution I google to is to add map_location='cpu'checkpoint = torch.load (PATH,map_location='cpu')

Load the part of the pre-training model:

Resnet152 = models.resnet152 (pretrained=True) # loading model structure and parameters pretrained_dict = resnet152.state_dict () "after loading the pre-training model and parameters in torchvision, the parameters can also be extracted by the state_dict () method directly from the official model_zoo: pretrained_dict = model_zoo.load_url (model_urls ['resnet152'])" model_dict = model.state_dict () # not in pretrained_dict The key that belongs to model_dict removes pretrained_dict = {k: v for k V in pretrained_dict.items () if k in model_dict} # Update the existing model_dictmodel_dict.update (pretrained_dict) # load the state_dictmodel.load_state_dict (model_dict) we really need

Or write in more detail:

Model_dict = model.state_dict () state_dict = {} for k, v in pretrained_dict.items (): if k in model_dict.keys (): # state_dict.setdefault (k V) state_ state_ [k] = velse:print ("Missing key (s) in state_dict: {}" .format (k)) model_dict.update (state_dict) model.load_state_dict (model_dict) the answer to the question about how the PyTorch model is saved and loaded is shared here I hope the above content can help you to a certain extent, if you still have a lot of doubts to be solved, you can follow the industry information channel to learn more related knowledge.

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.