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, task_size, task_rank, group=None, average=True)
| 1130 | allgather(x_list, torch.Tensor([x]).cuda()) |
| 1131 | |
| 1132 | def reduce_dict(input_dict, task_size, task_rank, group=None, average=True): |
| 1133 | """ |
| 1134 | Args: |
| 1135 | input_dict (dict): all the values will be reduced |
| 1136 | average (bool): whether to do average or sum |
| 1137 | Reduce the values in the dictionary from all processes so that all processes |
| 1138 | have the averaged results. Returns a dict with the same fields as |
| 1139 | input_dict, after reduction. |
| 1140 | """ |
| 1141 | world_size = task_size |
| 1142 | if world_size < 2: |
| 1143 | return input_dict |
| 1144 | with torch.no_grad(): |
| 1145 | names = [] |
| 1146 | values = [] |
| 1147 | # sort the keys so that they are consistent across processes |
| 1148 | for k in sorted(input_dict.keys()): |
| 1149 | names.append(k) |
| 1150 | values.append(input_dict[k]) |
| 1151 | values = torch.stack(values, dim=0) |
| 1152 | allreduce(values, group=group) |
| 1153 | if average: |
| 1154 | values /= world_size |
| 1155 | reduced_dict = {k: v for k, v in zip(names, values)} |
| 1156 | return reduced_dict |