将 Pytorch .pth 模型的权重保存到 .txt 或 .json

Question

我正在尝试将 pytorch 模型的权重保存到 .txt 或 .json 中。将其写入 .txt 时，

#import torch
model = torch.load("model_path")
string = str(model)
with open('some_file.txt', 'w') as fp:
     fp.write(string)

我得到一个文件，其中没有保存所有权重，即整个文本文件中都有省略号。我无法将其写入 JSON，因为该模型具有张量，这些张量不可 JSON 序列化 [除非有一种我不知道的方法？] 如何将权重保存在 .pth 文件中转换成某种格式，这样就不会丢失任何信息，并且可以很容易地看到？

谢谢

Answer 1

当你做str(model.state_dict())时，它递归地使用它包含的元素的str方法。所以问题是如何构建单个元素字符串表示。您应该增加以单个字符串表示形式打印的行数限制：

torch.set_printoptions(profile="full")

看看与此的区别：

import torch
import torchvision.models as models
mobilenet_v2 = models.mobilenet_v2()
torch.set_printoptions(profile="default")
print(mobilenet_v2.state_dict()['features.2.conv.0.0.weight'])
torch.set_printoptions(profile="full")
print(mobilenet_v2.state_dict()['features.2.conv.0.0.weight'])

张量目前 JSON 不可序列化。

将 Pytorch .pth 模型的权重保存到 .txt 或 .json

Saving the weights of a Pytorch .pth model into a .txt or .json

python

machine-learning

pytorch