如何腌制带有模型的字典？

Question

我在各种数据上拟合了变形金刚模型 DistillBert，并将每个拟合模型保存在字典中

{dataname : 拟合 DistillBert 模型}

如何用pickle保存这本字典？如果pickle不起作用，我该如何使用dill来保存这本字典并利用其中的模型？

Answer 1

如果您使用 pytorch，则无需自行使用开箱即用的 pickle 来保存模型。因为torch.save() 和torch.load() 负责处理它。 torch.save 函数使用 Python 的 pickle 实用程序进行序列化。您可以使用它保存模型。

在 Keras/Tensorflow 的情况下，您可以简单地使用 model.save() 而不是使用 pickle，即使 Keras 模型是 pickle-able。这样保存的模型将与未来版本的 Keras 兼容，并且可以在其他平台上使用。

在任何情况下，您都可以使用dill 来存储它：

import dill
...
...
#class definition of model1's class
...
#class definition of model2's class
...
d = dict(a1=model1, a2=model2)

with open("dillfilename", "wb") as dill_file:
    dill.dump(d , dill_file)

同样你可以使用dill.load来加载dill文件。