Spaces:
Running
Running
File size: 3,744 Bytes
15bcbe6 |
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 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 |
# Copyright 2022 Google.
#
# 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.
"""Helper routines for recording various training metrics."""
from typing import Any
import jax.numpy as jnp
Array = Any
def compute_accuracy_sum(logits, targets, valid_loss_mask=None):
"""Compute accuracy for logits and targets.
Args:
logits: [batch, length, num_classes] float array.
targets: categorical targets [batch, length] int array.
valid_loss_mask: None or array of shape bool[batch, length]
Returns:
The number of correct tokens in the output.
"""
if logits.shape[:-1] != targets.shape:
raise ValueError("Incorrect shapes. Got shape %s logits and %s targets" %
logits.shape, targets.shape)
if valid_loss_mask is not None and valid_loss_mask.shape != targets.shape:
raise ValueError("Incorrect shapes. Got shape %s targets and %s mask" %
targets.shape, valid_loss_mask.shape)
accuracy = jnp.equal(jnp.argmax(logits, axis=-1), targets)
if valid_loss_mask is not None:
accuracy = jnp.logical_and(accuracy, valid_loss_mask)
return jnp.sum(accuracy) # Sum of the number of True values.
def reshape_image(image):
"""Reshape image to something that tensorboard recognizes.
Args:
image: Array of shape [xsize, size] or [num_images, xsize, ysize]
Returns:
Array of shape [num_images, xsize, ysize, 1]
"""
# Reshape to [num_images, xdim, ydim, rgb] for tensorboard.
sh = image.shape
if image.ndim == 2:
return jnp.reshape(image, [1, sh[0], sh[1], 1]).astype(jnp.float32)
elif image.ndim == 3:
return jnp.reshape(image, [sh[0], sh[1], sh[2], 1]).astype(jnp.float32)
else:
return None # Not an image.
def normalize_image(images: Array, as_group: bool = False) -> Array:
"""Rescale the values in images to between 0.0 and 1.0.
Args:
images: Array of size [batch_size, xsize, ysize]
as_group: Scale all images in the batch by the same amount if True.
Returns:
A rescaled image of the same shape.
"""
images = images.astype(jnp.float32) # Return images as float32.
if as_group:
# Normalize the batch of images as a group.
min_img = jnp.min(images)
max_img = jnp.max(images)
else:
# Normalize each image in the batch individually.
min_img = jnp.min(images, axis=(-2, -1), keepdims=True)
max_img = jnp.max(images, axis=(-2, -1), keepdims=True)
norm_image = (images - min_img) / (max_img - min_img + 1e-6)
return jnp.where(jnp.isfinite(norm_image), norm_image, 0.0)
def overlay_images(image1: Array, image2: Array) -> Array:
"""Place image1 on top of image2, broadcasting image2 if necessary.
Args:
image1: array of shape [num_images, xsize, ysize]
image2: array of shape [num_images, xsize, ysize]
Returns:
A combined image.
"""
assert image1.ndim == 3 # (num_images, xsize, ysize)
assert image2.ndim == 3
image2 = jnp.broadcast_to(image2, image1.shape)
return jnp.concatenate([image1, image2], axis=1)
def make_histograms(viz_dicts):
"""Generate image histograms."""
hist_dict = {}
for (i, viz_dict) in enumerate(viz_dicts):
for (k, images) in viz_dict.items():
hist_dict["h_" + k + "_" + str(i)] = images
return hist_dict
|