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)
| 264 | |
| 265 | |
| 266 | def reduce_dict(input_dict, average=True): |
| 267 | """ |
| 268 | Args: |
| 269 | input_dict (dict): all the values will be reduced |
| 270 | average (bool): whether to do average or sum |
| 271 | Reduce the values in the dictionary from all processes so that all processes |
| 272 | have the averaged results. Returns a dict with the same fields as |
| 273 | input_dict, after reduction. |
| 274 | """ |
| 275 | world_size = get_world_size() |
| 276 | if world_size < 2: |
| 277 | return input_dict |
| 278 | with torch.no_grad(): |
| 279 | names = [] |
| 280 | values = [] |
| 281 | # sort the keys so that they are consistent across processes |
| 282 | for k in sorted(input_dict.keys()): |
| 283 | names.append(k) |
| 284 | values.append(input_dict[k]) |
| 285 | values = torch.stack(values, dim=0) |
| 286 | dist.all_reduce(values) |
| 287 | if average: |
| 288 | values /= world_size |
| 289 | reduced_dict = {k: v for k, v in zip(names, values)} |
| 290 | return reduced_dict |
| 291 | |
| 292 | |
| 293 | class MetricLogger(object): |
nothing calls this directly
no test coverage detected