pytorch保存 加载模型

state_dict

state_dict结构通过层去匹配参数张量

Define model

   class TheModelClass(nn.Module):
       def __init__(self):
           super(TheModelClass, self).__init__()
           self.conv1 = nn.Conv2d(3, 6, 5)
           self.pool = nn.MaxPool2d(2, 2)
           self.conv2 = nn.Conv2d(6, 16, 5)
           self.fc1 = nn.Linear(16 * 5 * 5, 120)
           self.fc2 = nn.Linear(120, 84)
           self.fc3 = nn.Linear(84, 10)

       def forward(self, x):
           x = self.pool(F.relu(self.conv1(x)))
           x = self.pool(F.relu(self.conv2(x)))
           x = x.view(-1, 16 * 5 * 5)
           x = F.relu(self.fc1(x))
           x = F.relu(self.fc2(x))
           x = self.fc3(x)
           return x

Initialize model

   model = TheModelClass()

Initialize optimizer

   optimizer = optim.SGD(model.parameters(), lr=0.001, momentum=0.9)

Print model’s state_dict

   print("Model's state_dict:")
   for param_tensor in model.state_dict():
       print(param_tensor, "\t", model.state_dict()[param_tensor].size())

Print optimizer’s state_dict

   print("Optimizer's state_dict:")
   for var_name in optimizer.state_dict():
       print(var_name, "\t", optimizer.state_dict()[var_name])

Output:

::

   Model's state_dict:
   conv1.weight     torch.Size([6, 3, 5, 5])
   conv1.bias   torch.Size([6])
   conv2.weight     torch.Size([16, 6, 5, 5])
   conv2.bias   torch.Size([16])
   fc1.weight   torch.Size([120, 400])
   fc1.bias     torch.Size([120])
   fc2.weight   torch.Size([84, 120])
   fc2.bias     torch.Size([84])
   fc3.weight   torch.Size([10, 84])
   fc3.bias     torch.Size([10])

   Optimizer's state_dict:
   state    {}
   param_groups     [{'lr': 0.001, 'momentum': 0.9, 'dampening': 0, 'weight_decay': 0, 'nesterov': False, 'params': [4675713712, 4675713784, 4675714000, 4675714072, 4675714216, 4675714288, 4675714432, 4675714504, 4675714648, 4675714720]}]

state_dict格式的model的保存与加载

torch.save(model.state_dict(), PATH)#保存

model = TheModelClass(*args, **kwargs)#加载
model.load_state_dict(torch.load(PATH))
model.eval()

检查点的模型保存与加载（一个模型的多种属性

torch.save({
               'epoch': epoch,
               'model_state_dict': model.state_dict(),
               'optimizer_state_dict': optimizer.state_dict(),
               'loss': loss,
               ...
               }, PATH)#保存

model = TheModelClass(*args, **kwargs)
optimizer = TheOptimizerClass(*args, **kwargs)
checkpoint = torch.load(PATH)   model.load_state_dict(checkpoint['model_state_dict'])
optimizer.load_state_dict(checkpoint['optimizer_state_dict'])
epoch = checkpoint['epoch']
loss = checkpoint['loss']
model.eval()
   # - or -
model.train()

保存或加载多个模型

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 = TheModelBClass(*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()
   # - or -
   modelA.train()
   modelB.train()