Multi-GPU Inference

[angadkalra]

Is it possible to do DDP multi-gpu inference and write to a file at end? I tried and I only get the outputs of one GPU in the file. Is there a way to all gather at end of test epoch and write all the output_dicts returned at end of every step to a single .pkl file?

[tchaton]

Dear @angadkalra,

You could rely on all_gather to send the batches across, but there aren't the best practices as it might be costly and quite error-prone.
I would recommend using the PredictionWriter Callback to write batches of predictions to the disk and then use rank 0 to merge all predictions together.

An even better solution would be to write down the predictions to a database with a file lock for each process to write on.

[lucifermorningstar1305]

Hi @tchaton,

If it's possible, can you please share an example?