def masked_mean(inputs, targets, mask_id=None):
"""Mean of the inputs but counting only those where targets != mask_id."""
x = inputs.astype(np.float32)
if mask_id is None:
return np.mean(x)
unmask = 1.0 - np.equal(targets, mask_id).astype(np.float32)
return np.sum(x * unmask) / np.sum(unmask)
def accuracy(batch, model_predictions):
"""Calculate accuracy."""
_, targets = batch
model_predictions, targets = _make_list(model_predictions, targets)
correct = []
for (prediction, target) in zip(model_predictions, targets):
predicted_class = np.argmax(prediction, axis=-1)
correct.append(np.equal(predicted_class, target))
return masked_mean(correct, targets)
def masked_mean(inputs, targets, mask_id=None):
"""Mean of the inputs but counting only those where targets != mask_id."""
inputs = [x.astype(np.float32) for x in inputs]
# We assume all elements in the list contribute equally.
# TODO(lukaszkaiser): remove this assumption (e.g., when masks differ).
length = len(inputs)
if mask_id is None:
# TODO(lukaszkaiser): can we just divide the sum by length? XLA optimizes?
return sum([np.mean(x) / length for x in inputs])
unmask = [1.0 - np.equal(t, mask_id).astype(np.float32) for t in targets]
return sum([np.sum(x * m) / (length * np.sum(m))
for x, m in zip(inputs, unmask)])
def masked_mean(inputs, targets, weights, mask_id=None):
"""Weighted mean of the inputs, excluding where targets == mask_id."""
inputs = [x.astype(np.float32) for x in inputs]
# We assume all elements in the list contribute equally.
# TODO(lukaszkaiser): remove this assumption (e.g., when masks differ).
length = len(inputs)
if mask_id is not None:
weights = [w * (1.0 - np.equal(t, mask_id).astype(np.float32))
for t, w in zip(targets, weights)]
weight_sums = [np.float32(t.size) if np.isscalar(w) else np.sum(w)
for w, t in zip(weights, targets)]
return sum([np.sum(x * w) / (length * s)
for x, w, s in zip(inputs, weights, weight_sums)])
def accuracy(batch, model_predictions): """Calculate accuracy.""" _, targets = batch predicted_class = np.argmax(model_predictions, axis=-1) correct = np.equal(predicted_class, targets) return masked_mean(correct, targets)
def WeightMask(target, mask_id=0, **kw):
del kw
if mask_id is None:
return np.ones_like(target)
return 1.0 - np.equal(target, mask_id).astype(np.float32)
def Accuracy(x, axis=-1, **kw):
del kw
prediction, target = x
predicted_class = np.argmax(prediction, axis=axis)
return np.equal(predicted_class, target)
