pytorch之保存与加载模型

本篇笔记译自pytorch官网tutorial，用于方便查看。
pytorch与保存、加载模型有关的常用函数3个：

• torch.save(): 保存一个序列化的对象到磁盘，使用的是Python的pickle库来实现的。
• torch.load(): 解序列化一个pickled对象并加载到内存当中。
• torch.nn.Module.load_state_dict(): 加载一个解序列化的state_dict对象

1. state_dict

在PyTorch中所有可学习的参数保存在model.parameters()中。state_dict是一个Python字典。保存了各层与其参数张量之间的映射。torch.optim对象也有一个state_dict，它包含了optimizer的state，以及一些超参数。

2. 保存&加载模型来inference(recommended)

save

torch.save(model.state_dict(), PATH)

load

model = TheModelClass(*args, **kwargs)
model.load_state_dict(torch.load(PATH))
model.eval()  # 当用于inference时不要忘记添加

• 保存的文件名后缀可以是.pt或.pth
• 当用于inference时不要忘记添加model.eval()

3. 保存&加载整个模型(not recommended)

save

torch.save(model, PATH)

load

# Model class must be defined somewhere
model = torch.load()
model.eval()

4. 保存&加载带checkpoint的模型用于inference或resuming training

save

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

load

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()

5. 保存多个模型到一个文件中

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)

load

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

• 此情况可能在GAN，Sequence-to-sequence，或ensemble models中使用
• 保存checkpoint常用.tar文件扩展名

6. Warmstarting Model Using Parameters From A Different Model

save

torch.save(modelA.state_dict(), PATH)

load

modelB = TheModelBClass(*args, **kwargs)
modelB.load_state_dict(torch.load(PATH), strict=False)

• 在迁移训练时，可能希望只加载部分模型参数，此时可置strict参数为False来忽略那些没有匹配到的keys

7. 保存&加载模型跨设备

(1) Save on GPU, Load on CPU
save

torch.save(model.state_dict(), PATH)

load

device = torch.device("cpu")
model = TheModelClass(*args, **kwargs)
model.load_state_dict(torch.load(PATH, map_location=device))

(2) Save on GPU, Load on GPU
save

torch.save(model.state_dict(), PATH)

load

device = torch.device("cuda")
model = TheModelClass(*args, **kwargs)
model.load_state_dict(torch.load(PATH))
model.to(device)

(3) Save on CPU, Load on GPU
save

torch.save(model.state_dict(), PATH)

load

device = torch.device("cuda")
model = TheModelClass(*args, **kwargs)
model.load_state_dict(torch.load(PATH, map_location="cuda:0"))
model.to(device)

8. 保存torch.nn.DataParallel模型

save

torch.save(model.module.state_dict(), PATH)

load

# Load to whatever device you want

---