问题描述
是否有约定在 PyTorch Lightning 中实现某种 predict()
方法,在使用 forward()
执行实际预测之前进行预处理?
就我而言,我有一个文本分类器,由一个嵌入层和几个全连接层组成。文本在传递到嵌入层之前需要进行标记。在训练和评估期间,LightningDataModule
的 setup()
方法可以完成这项工作。
现在,我想知道在生产过程中进行推理的最佳实践是什么。我可以将 predict()
方法添加到我的 LightningModule
中,我可以在其中编写与 LightningDataModule.setup()
中相同的预处理代码。但是,当然,我不想复制代码。
在 PyTorch Lightning 官方文档中链接的 this community example 项目中,作者在 prepare_sample()
中定义了一个 LightningModule
函数,该函数被他们的 predict()
函数使用,并且也是传递给 LightningDataModule
。
这是处理预处理的正确方法吗?另外,为什么 LightningModule
中没有 prepare_sample()
或 predict()
?对我来说,这似乎是一个常见的用例,例如:
model = load_model('data/model.ckpt') # load pre-trained model,analyzes user reviews
user_input = input('Your movie review > ')
predicted_rating = model.predict(user_input) # e.g. "I liked the movie pretty much." -> 4 stars
print('Predicted rating: %s/5 stars' % predicted_rating)
现在我想起来了,predict()
也应该像评估代码一样处理 forward()
的结果,比如选择输出最高的类或选择所有输出大于一些阈值 - 一些不应重复的代码。
解决方法
如果代码应该用于生产,为什么要使用 LightningModule
?如果模型完成,您只需从内存中加载模型并定义预处理步骤。
您引用的存储库在 LightningModule 之上实现了 predict 和 prepare_sample。
在我看来,pytorch-lightning 是用于模型的训练和评估,而不是用于生产。我们不希望在将模型发送到生产环境时保留分析和调试,因此我们创建了一个精简版,只加载模型、预处理和预测。
Towardsdatascience 有一个小代码示例:https://towardsdatascience.com/how-to-deploy-pytorch-lightning-models-to-production-7e887d69109f