File size: 3,070 Bytes
74e8f2f |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 |
# Copyright 2022 Big Vision Authors.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
"""Evaluator for computing mean of per-example metrics."""
import functools
from typing import Mapping
from big_vision import input_pipeline
from big_vision.datasets import core as ds_core
from big_vision.pp import builder as pp_builder
import jax
import jax.numpy as jnp
import numpy as np
# Note: global to avoid jax re-compiling across different evaluator instances.
@functools.partial(jax.pmap, static_broadcasted_argnums=0, axis_name='batch')
def _run_predict_fn(predict_fn, params, batch):
"""Sum per-example metrics weighted by `_mask`."""
mask = batch['_mask']
metrics = predict_fn(params, batch)
# Sanity check output format of predict_fn.
assert isinstance(metrics, Mapping), 'predict_fn must return a dict'
for y in jax.tree_leaves(metrics):
if y.shape != mask.shape:
raise ValueError(
f'Expected per-example metrics of shape {mask.shape} found '
f'{jax.tree_map(lambda x: x.shape, metrics)}.')
metrics = {**metrics, '_mask': mask}
metrics = jax.tree_map(lambda x: jnp.inner(x, mask), metrics)
return jax.lax.psum(metrics, axis_name='batch')
class Evaluator:
"""Report the mean of per-example metrics computed by predict_fn.
`predict_fn(params, batch)` must return a dict from metric name to
per-example metrics of shape [batch_size].
"""
def __init__(self, predict_fn, data, pp_fn, batch_size,
cache_final=True, cache_raw=False, prefetch=1):
data = ds_core.get(**data)
self.dataset, self.steps = input_pipeline.make_for_inference(
data.get_tfdata(ordered=True), batch_size=batch_size,
num_ex_per_process=data.num_examples_per_process(),
preprocess_fn=pp_builder.get_preprocess_fn(pp_fn),
cache_final=cache_final, cache_raw=cache_raw)
self.data_iter = input_pipeline.start_input_pipeline(self.dataset, prefetch)
self.predict_fn = predict_fn
def run(self, params):
"""Computes all metrics."""
metrics = []
# Compute batch metrics without blocking.
for _, batch in zip(range(self.steps), self.data_iter):
batch_metrics = _run_predict_fn(self.predict_fn, params, batch)
metrics.append(batch_metrics)
# Transfer metrics from device 0 to host (blocking).
metrics = jax.device_get(jax.tree_map(lambda x: x[0], metrics))
metrics_sum = jax.tree_map(lambda *x: np.sum(x), *metrics)
mask_sum = metrics_sum.pop('_mask')
for key, value_sum in metrics_sum.items():
yield (key, value_sum / mask_sum)
|