使用pytorch照明时获取模型输出

Question

我正在为一个项目使用pytorch lighting。我的模型输出单词嵌入，我想把模型输出保存在一个文件中。培训完成后，我如何访问模型输出？我应该使用回调吗？在使用training_end_epoch((时，当我检查train_outs时，我只得到一个批次的输出，而不是所有批次的输出。

Fyi，由于一些代码限制，我使用pytorch lighting 0.7.6。

Answer 1

我假设您从training_step返回单词embeddings。假设您返回一个字典{'embed': var_embed}。

您可以在training_epoch_end (self, outputs)中访问该词典。您可以使用列表收集所有批次的输出。比如

embeds = [x['embed'] for x in outputs]

请发布MWE，这会更容易。