Could `join` replace `gather_for_metrics` and perform it automatically?

Hi,

Great job with accelerate!

One persistent headache that I keep experiencing is the repeated samples during distributed evaluation. Although there is the gather_for_metrics functionality this doesn’t always work depending on the output of the model; for example, Faster-RCNN in Torchvision, which outputs a list of dictionaries containing tensors.

I was wondering if it would be possible to remove the repeated behaviour completely, using something like the join context manager, which enables training on uneven outputs.

If there is appetite for this, I would be happy to help you explore options.

Thanks

[sgugger]

Accelerate was built when there was no join contextmanager and the fact each datalaoder returns the same number of samples in all processes is pretty ingrained in the library. We could make this evolve in the future with a major release, but that would mean a lot of changes (and breaking functionality).

For now I’d rather investigate why your use case is not supported as is (all methods should support a list of dictionaries) rather than rewrite the whole library.

[sgugger]

@Chris-hughes10 This is mostly used in the whole evaluation part of Accelerate (so gather and gather_for_metrics) as for training we don’t really care (plus the dataloaders very often have drop_last=True during training so there is no problem there).

I’m open to start exploring a different way to go, probably with a new flag in the accelerator. The first thing would be to have the batch sampler we have handle a non-fixed batch size (which would also be useful for training), then add this flag where we would not cycle through the dataset but return different lengths on different processes and finally add a wrapper around join.

Does that sound reasonable? If so, I can start a project summarizing the steps and we can share the work.