Batches with Tags

[imagirom]

I think it would be useful to have a nicer interface for input and output batches with multiple elements.

Currently it can be cumbersome to keep track of what is where in the batch, especially when using multiple transforms that add or remove elements from the batch, or when using multiple loss functions that act on different ground truth/predictions.

This would be a lot easier if elements in the batch could have tags (such as 'raw', 'segmentation', 'affinities'). Transforms and loss functions could use these tags to select what they act on, and also label their outputs.

I will probably implement this at least for myself, but doing so in a nice way while keeping the current functionality will be harder. So I am interested in whether this feature would be useful to others, and if someone has ideas on how to implement it.