pytorch-lightning - validation_epoch_end 与 DDP Pytorch Lightning

Question

你的问题是什么？

我正在尝试实现一个需要访问整个数据的指标。因此，我没有在 *_step() 方法中更新指标，而是尝试在 *_epoch_end() 方法中收集输出。但是，输出仅包含每个设备获取的数据分区的输出。基本上，如果有 n 个设备，那么每个设备都会获得总输出的 1/n。

你的环境是什么？

OS: ubuntu
Packaging: conda
Version [1.0.4
Pytorch: 1.6.0

Answer 1

参见 pytorch-lightningmanual .我认为您正在寻找 training_step_end/validation_step_end(假设您使用的是 DP/DDP2)。

...So, when Lightning calls any of the training_step, validation_step, test_step you will only be operating on one of those pieces. (...) For most metrics, this doesn’t really matter. However, if you want to add something to your computational graph (like softmax) using all batch parts you can use the training_step_end step.