Add batched processing and augmentation for N2N

[Obi-Wan]

Description

Noise2Noise can be used with large stacks of images, and they might not fit into GPU memory. Some form of batching would be useful for accumulating the gradients over a large amount of data.

Moreover, data augmentation techniques could also improve the learning quality, without requiring additional data. Common data augmentation techniques include image or volume flips and rotations.

TODO

• Implement the feature:
  • Implement batching along the tensor designated batch dimension
  • Implement data augmentation techniques (at least the flips)
  • Consider where to always do gradient accumulation or to add support for batched iterations
• Update docstrings
• Update / add tests
• Update documentation

Notes

These two features complement each other well because they both reduce the memory requirements needed to obtain a high-quality model.

[Obi-Wan]

In the end, we opted for having the choice between gradient accumulation and per-batch updates, depending on what the user prefers (switched on/off with the parameter accum_grads).