pytorch模型文件pth详解

Posted 2023-03-04

参考技术A

如上打印输出所示，pth文件通过有序字典来保持模型参数。有序字典与常规字典一样，但是在排序操作方面有一些额外的功能。常规的dict是无序的，OrderedDict能够比dict更好地处理频繁的重新排序操作。
OrderedDict有一个方法 popitem(last=True) 用于有序字典的popitem()方法返回并删除一个(键，值)对。如果last为真，则按LIFO顺序返回对;如果为假，则按FIFO顺序返回对。
OrderedDict还有一个方法 move_to_end(key,last=True) ,将现有的键移动到有序字典的两端。如果last为真，则将项目移动到右端(默认);如果last为假，则移动到开头。

如上打印所示，有序字典state_dict中每个元素都是Parameter参数，该参数是一种特殊的张量，包含data和requires_grad两个方法。其中data字段保存的是模型参数，requires_grad字段表示当前参数是否需要进行反向传播。

先建立一个字典，保存三个参数：调用torch.save(),即可保存对应的pth文件。需要注意的是若模型是由nn.Moudle类继承的模型，保存pth文件时，state_dict参数需要由 model.state_dict 指定。

当你想恢复某一阶段的训练（或者进行测试）时，那么就可以读取之前保存的网络模型参数等。