Pytorch模型存储心得小记

本文持续更新使用Pytorch存储模型的心得小记。

1. 推荐torch.save(model.state_dict(), ‘dir’)而不是torch.save(model, ‘dir’)

二者都可以用torch.load()重新加载，不同在于model是整个模型，model.state_dict()只含模型参数，我们常下载使用的预训练模型的checkpoint就是state_dict。

实际使用时，用torch.save(model, ‘dir’)方便，配合 model=torch.load('dir') 即可重新加载这个模型；而torch.save(model.state_dict(), ‘dir’)稍稍麻烦些，配合 import BertModel 、model=BertModel.from_pretrain("bert-base-uncased") 、model.load_state_dict('dir') 这3条语句才可重新加载模型。

既然torch.save(model, ‘dir’)使用起来更方便，为什么不推荐呢？原因在于：模型版本一旦发生变化，我们极可能无法重新加载存储好的model。以我的经历举例，从transformers 4.3.3升级为4.15.0，原先torch.save的model，都不能torch.load了，有error报错，比如‘BertEncoder’ object has no attribute ‘gradient_checkpointing’。此外，也建议用save_pretrained/from_pretrained。

2. model.state_dict()和model.named_parameters()并不仅仅是数据类型不同

二者最直观的不同在于，state_dict是OrderdDict类型，配合state_dict()[参数名]就能取得参数值，named_parameters是generator类型，只能for循环或者.next()依次读取参数值。

二者的另一个不同在于，参数项可能不同。还是以我的经历举例，t5模型的named_parameters()的参数比state_dict()的参数少了两项： encoder.embed_tokens.weight 和 decoder.embed_tokens.weight。需要注意的是，新模型加载参数时会 随机初始化这两个缺少的参数项 ，模型效果会有所下降。所以存储模型参数时，要存储state_dict，而不是named_parameters。