Function reduce_dict

util/misc.py:157–194

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)

def reduce_dict(input_dict, average=True):
    """
    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 reduction.
    """
    world_size = get_world_size()
    if world_size < 2:
        return input_dict
    with torch.no_grad():
        names = []
        values = []
        # sort the keys so that they are consistent across processes
        # import pdb; pdb.set_trace()
        for k in sorted(input_dict.keys()):

            names.append(k)
            values.append(input_dict[k])
        # pdb.set_trace()
        values = torch.stack(values, dim=0)
        
        try:
            dist.all_reduce(values)
            rank = dist.get_rank()
            # logging.info(f'Rank {rank} after all_reduce')
        except Exception as e:
            rank = dist.get_rank()
            print(f'Exception in rank {rank}: {e}')
            # print(f'values: {values}')
            # print(f'names: {names}')
            logging.info(f'Rank {rank} after all_reduce')
        if average:
            values /= world_size
        reduced_dict = {k: v for k, v in zip(names, values)}
    return reduced_dict

def setup_logging():
    logging.basicConfig(level=logging.INFO)