Args: input_dict (dict): all the values will be reduced average (bool): whether to do average or sum Reduce the values in the dictionary from all processes so that all processes have the averaged results. Returns a dict with the same fields as input_dict, after reduc
(input_dict, average=True)
| 84 | |
| 85 | # Function from DETR - https://github.com/facebookresearch/detr/blob/master/util/misc.py |
| 86 | def reduce_dict(input_dict, average=True): |
| 87 | """ |
| 88 | Args: |
| 89 | input_dict (dict): all the values will be reduced |
| 90 | average (bool): whether to do average or sum |
| 91 | Reduce the values in the dictionary from all processes so that all processes |
| 92 | have the averaged results. Returns a dict with the same fields as |
| 93 | input_dict, after reduction. |
| 94 | """ |
| 95 | world_size = get_world_size() |
| 96 | if world_size < 2: |
| 97 | return input_dict |
| 98 | with torch.no_grad(): |
| 99 | names = [] |
| 100 | values = [] |
| 101 | # sort the keys so that they are consistent across processes |
| 102 | for k in sorted(input_dict.keys()): |
| 103 | names.append(k) |
| 104 | values.append(input_dict[k]) |
| 105 | values = torch.stack(values, dim=0) |
| 106 | torch.distributed.all_reduce(values) |
| 107 | if average: |
| 108 | values /= world_size |
| 109 | reduced_dict = {k: v for k, v in zip(names, values)} |
| 110 | return reduced_dict |
| 111 | |
| 112 | |
| 113 | # Function from https://github.com/facebookresearch/detr/blob/master/util/misc.py |
nothing calls this directly
no test coverage detected