def compute_multi_gpu_topk_accuracy()

in lib/utils/metrics.py [0:0]


def compute_multi_gpu_topk_accuracy(top_k, split):

    aggr_batch_size = 0
    aggr_top_k_correct_hits = 0

    if cfg.METRICS.EVAL_FIRST_N and split in ['test', 'val']:
        aggr_top_k_N_way_correct_hits = 0

    computed_metrics = {}

    for idx in range(cfg.ROOT_GPU_ID, cfg.ROOT_GPU_ID + cfg.NUM_GPUS):

        softmax = workspace.FetchBlob('gpu_{}/pred'.format(idx))
        # remove the last two dimensions if we use conv as the output fc
        softmax = softmax.reshape((softmax.shape[0], -1))
        labels = workspace.FetchBlob('gpu_{}/labels'.format(idx))

        assert labels.shape[0] == softmax.shape[0], "Batch size mismatch."

        aggr_batch_size += labels.shape[0]
        aggr_top_k_correct_hits += compute_topk_correct_hits(
            top_k, softmax, labels)

        if cfg.METRICS.EVAL_FIRST_N and split in ['test', 'val']:
            aggr_top_k_N_way_correct_hits += compute_topk_correct_hits(
                top_k, softmax[:, :cfg.METRICS.FIRST_N], labels)

    # normalize results
    computed_metrics['topk_accuracy'] = \
        float(aggr_top_k_correct_hits) / aggr_batch_size
    if cfg.METRICS.EVAL_FIRST_N and split in ['test', 'val']:
        computed_metrics['topk_N_way_accuracy'] = \
            float(aggr_top_k_N_way_correct_hits) / aggr_batch_size

    return computed_metrics