CODEWITHSUNDEEP
pythonimportneural-networkpytorchexport
Leo
Leo

Reputation: 781

Pytorch creating model from load_state_dict

I'm trying to learn how to save and load trained models in Pytorch, but so far, I'm only getting errors. Let's consider the following self-contained code:

import torch
lin=torch.nn.Linear; act=torch.nn.ReLU(); fnc=torch.nn.functional;
class Ann(torch.nn.Module): 
   def __init__(self):
      super(Ann, self).__init__() 
      self.conv1 = torch.nn.Conv2d( 1, 10, kernel_size=5) 
      self.conv2 = torch.nn.Conv2d(10, 20, kernel_size=4) 
      self.drop = torch.nn.Dropout2d(p=0.5) 
      self.fc1 = torch.nn.Linear(320,128)  
      self.fc2 = torch.nn.Linear(128,10)   
   def forward(self, x): 
      x = self.conv1(x[:,None,:,:]);        
      x = fnc.relu(fnc.max_pool2d(x,2));    
      x = self.drop(self.conv2(x));         
      x = fnc.relu(fnc.max_pool2d(x,2));    
      x = torch.flatten(x,1);               
      x = fnc.relu(self.fc1(x));            
      x = fnc.dropout(self.fc2(x),training=self.training);   
      return fnc.log_softmax(x,dim=0) 
x,y=torch.rand((5,28,28)),torch.randint(0,9,(5,)); 
f=fnc.nll_loss;   
ann1 = torch.nn.Sequential( torch.nn.Flatten(start_dim=1), 
   lin(784,256), act, lin(256,128), act, lin(128,10), torch.nn.LogSoftmax(dim=1)) 
ann2=Ann()
F1 = torch.optim.SGD(ann1.parameters(),lr=0.01,momentum=0.5)
F2 = torch.optim.SGD(ann2.parameters(),lr=0.01,momentum=0.5)

F1.zero_grad(); y_=ann1(x); loss=f(y_,y); loss.backward(); F1.step()
print(x.dtype,y.dtype,x.shape,y.shape,y_.shape,loss); 
F2.zero_grad(); y_=ann2(x); loss=f(y_,y); loss.backward(); F2.step()
print(x.dtype,y.dtype,x.shape,y.shape,y_.shape,loss); 
name='/home/leon/'

#ann3 = ann1.__class__().load_state_dict(ann1.state_dict());  print(ann3(x)) #outputs errors
#ann4 = ann2.__class__().load_state_dict(ann2.state_dict());  print(ann4(x)) #outputs errors
torch.save( [ann1.state_dict(),F1.state_dict()], name+'annF1.pth'); 
torch.save( [ann2.state_dict(),F2.state_dict()], name+'annF2.pth'); 
a1,d1=torch.load(name+'annF1.pth')
a2,d2=torch.load(name+'annF2.pth') #so far, works as expected
ann3, F3 = ann1.__class__().load_state_dict(a1), F1.__class__().load_state_dict(d1) #outputs errors
ann4, F4 = ann2.__class__().load_state_dict(a2), F2.__class__().load_state_dict(d2) #outputs errors

As you can see, ann1 and ann2 work, since they produce valid output. However, (re)constructing a model ann3 and ann4 from the given state_dict() invariably gives two errors (respectively):

Unexpected key(s) in state_dict: "1.weight", "1.bias", "3.weight", "3.bias", "5.weight", "5.bias". 
TypeError: '_IncompatibleKeys' object is not callable

Could anyone please show me how to properly construct a model from given parameters, so I can later export and import my trained models?

Upvotes: 1

Views: 5819

Answers (1)

Niv Dudovitch
Niv Dudovitch

Reputation: 1658

Hey you have two problems:

  1. Remove the .__class__()
  2. Separate the definition of ann3 and ann4.
ann1.load_state_dict(ann1.state_dict())
ann3 = ann1
print(ann3(x))
ann2.load_state_dict(ann2.state_dict())
ann4 = ann2
print(ann4(x))

But, what is the propose of thisann1.__class__().load_state_dict(ann1.state_dict())?

Maybe you wanted to do this?

ann3 = torch.nn.Sequential( torch.nn.Flatten(start_dim=1), 
   lin(784,256), act, lin(256,128), act, lin(128,10), torch.nn.LogSoftmax(dim=1))
ann3.load_state_dict(ann1.state_dict())
print(ann3(x))

ann4 = Ann()
ann4.load_state_dict(ann2.state_dict())
print(ann4(x))

Its Works the same as the guide here, creates a new model with the same architecture, and then loads the saved/exist state_dict. Saving & Loading Model for Inference

model = TheModelClass(*args, **kwargs)
model.load_state_dict(torch.load(PATH))

Upvotes: 3

Related Questions