# HF Trainer downstream evaluation on multiple GPUS

**URL:** </static-proxy?url=https%3A%2F%2Fdiscuss.huggingface.co%2Ft%2Fhf-trainer-downstream-evaluation-on-multiple-gpus%2F27971%3E%5C
**Category:** 🤗Transformers\
**Created:** [December 16, 2022, 8:23pm UTC](/static-proxy?url=https%3A%2F%2Fdiscuss.huggingface.co%2Ft%2Fhf-trainer-downstream-evaluation-on-multiple-gpus%2F27971 "2022-12-16T20:23:34Z")\
**Posts on this page:** 2\
**Page:** 1

<div class="post-metadata">

**Author:** ![dhansmair](https://sea2.discourse-cdn.com/hellohellohello/user_avatar/discuss.huggingface.co/dhansmair/32/13081_2.png) [@dhansmair](/static-proxy?url=https%3A%2F%2Fdiscuss.huggingface.co%2Fu%2Fdhansmair)%5C
**Post date:** [December 16, 2022, 8:23pm UTC](/static-proxy?url=https%3A%2F%2Fdiscuss.huggingface.co%2Ft%2Fhf-trainer-downstream-evaluation-on-multiple-gpus%2F27971%2F1 "2022-12-16T20:23:34Z")

</div>

Hi all, I’m trying to train a language model using HF Trainer on four GPUs (multi-GPU newbie here). During evaluation, I want to track performance on downstream tasks, e.g. Image Captioning on COCO. I have overridden the evaluate() method and created the evaluation dataset in it. From the logs I can see that now during training, evaluation runs on all four GPUs in parallel, but on the same data.  
My question is, is there a simple way to split the dataset into four subsets, let each GPU run prediction on a subset, and then merge the results in the end? My initial idea was to split the dataset based on the rank of the current process, but how can I send the results e.g. to the master process (rank 0)? This process would then compute the metrics for all predictions.

Thank you a lot! 🤗

---

<div class="post-metadata">

**Author:** ![dhansmair](https://sea2.discourse-cdn.com/hellohellohello/user_avatar/discuss.huggingface.co/dhansmair/32/13081_2.png) [@dhansmair](/static-proxy?url=https%3A%2F%2Fdiscuss.huggingface.co%2Fu%2Fdhansmair)%5C
**Post date:** [December 21, 2022, 9:52am UTC](/static-proxy?url=https%3A%2F%2Fdiscuss.huggingface.co%2Ft%2Fhf-trainer-downstream-evaluation-on-multiple-gpus%2F27971%2F2 "2022-12-21T09:52:49Z")

</div>

I ended up using a small helper function like this:

```python

def evaluate_distributed(
    self, 
    dataset: Dataset, 
    eval_func: Callable[[Dataset], Dict[str, float]]
) -> Dict[str, float]:

    # split into even subsets based on the local rank of each GPU
    n = len(dataset) // self.args.world_size
    
    subset = Subset(
        dataset, 
        range(self.args.local_rank*n, (self.args.local_rank+1)*n)
    )
    
    metrics = eval_func(subset)
    gathered_metrics = [None for _ in range(self.args.world_size)]
    dist.all_gather_object(gathered_metrics, metrics)

    all_metrics = {}
    for k, v in metrics.items():
        l = [d[k] for d in gathered_metrics]
        if isinstance(v, float):
            all_metrics[k] = sum(l) / len(l)
        else:
            all_metrics[k] = l

    return all_metrics

```

This method takes a dataset and a function that computes a metric dictionary from a dataset.  
`dist.all_gather_object()` will distribute the metrics dict across all distributed processes.  
Of course, it needs to be valid to compute the final metric as the mean of the metric on the individual splits.