简体   繁体   中英

How can I load and use a PyTorch (.pth.tar) model

I am not very familiar with Torch, and I primarily use Tensorflow. I, however, need to use a retrained inception model that was retrained in Torch. Due to the large amount of computing resources required to retrain an inception model for my particular application, I would like to use the model that was already retrained.

This model is saved as a .pth.tar file.

I would like to be able to first load this model. So far, I have been able to figure out that I must use the following:

model = torch.load('iNat_2018_InceptionV3.pth.tar', map_location='cpu')

This seems to work, because print(model) prints out a large set of numbers and other values, which I presume are the values for the weights an biases.

After this, I need to be able to classify an image with it. I haven't been able to figure this out. How must I format the image? Should the image be converted into an array? After this, how must I pass the input data to the network?

you basically need to do the same as in tensorflow. That is, when you store a network, only the parameters (ie the trainable objects in your network) will be stored, but not the "glue", that is all the logic you need to use a trained model. So if you have a .pth.tar file, you can load it, thereby overriding the parameter values of a model already defined.

That means that the general procedure of saving/loading a model is as follows:

  • write your network definition (ie your nn.Module object)
  • train or otherwise change the network's parameters in a way you want
  • save the parameters using torch.save
  • when you want to use that network, use the same definition of an nn.Module object to first instantiate a pytorch network
  • then override the values of the network's parameters using torch.load

Here's a discussion with some references on how to do this: pytorch forums

And here's a super short mwe:

# to store
torch.save({
    'state_dict': model.state_dict(),
    'optimizer' : optimizer.state_dict(),
}, 'filename.pth.tar')

# to load
checkpoint = torch.load('filename.pth.tar')
model.load_state_dict(checkpoint['state_dict'])
optimizer.load_state_dict(checkpoint['optimizer'])

The technical post webpages of this site follow the CC BY-SA 4.0 protocol. If you need to reprint, please indicate the site URL or the original address.Any question please contact:yoyou2525@163.com.

 
粤ICP备18138465号  © 2020-2024 STACKOOM.COM