At some point I’ve looked into Stochastic Weight Averaging 63, which claims that a simple averaging of multiple checkpoints leads to a better generalization.
If I’m not mistaken, there is also an SWA package for PyTorch, which applies this strategy during your training.
That being said, I don’t know if these experiments still hold true for other models than the one mentioned in the paper.
Generally, I would assume that this method would only work, if all training checkpoints are close to the global minimum, which seems to be the case if you use the checkpoints from a single run (if I remember the paper correctly).
@omarfoq I think this operation might benefit from the current ongoing port of apex' multi_tensor_apply in this PR 22.
You could have a look at the code changes and check, if you could reuse it.
However, as mentioned before, I would profile the code beforehand and make sure that a speedup would be visible in your case.