Losses

The losses of torch.nn.functional, one value a row: a row of a [R, C] tile is one sample, and each loss leaves its [R, 1] column, a sample’s loss. PyTorch reduces over the batch as well, with reduction="mean" or "sum"; a kernel does that itself, often together with the other blocks’ rows, so the library stops at the row.

kernel token_loss(Logits: tensor<f32>[M, 512], Target: tensor<i32>[M, 1], Loss: tensor<f32>[M, 1]) {
  let r = ctx.block.x * 8
  var x = Logits[r :+ 8, :]
  cross_entropy(Loss[r :+ 8, :], x, Target[r :+ 8, :])
}

A target is a column of class indices, i32. The logit at each row’s class is picked by a where against iota(1) and a row sum, so no thread reads a column it does not already hold.

Classification

cross_entropy

fn cross_entropy(out, t, target)statement

Cross-entropy of the logits t ([R, C]) against the class of each row, target ([R, 1], i32): logsumexp(t) - t[target], into out ([R, 1]). t is left unchanged.

Source, std/loss.ph:9

fn cross_entropy(out, t, target) {
  logsumexp(out, t)
  out = out - rowsum(where(iota(1) == target, t, 0.0))
}

nll_loss

fn nll_loss(logp, target)expression

The negative log-likelihood of each row's class target ([R, 1], i32) under the log-probabilities logp, as log_softmax leaves them: -logp[target].

Source, std/loss.ph:17

fn nll_loss(logp, target) = -rowsum(where(iota(1) == target, logp, 0.0))

binary_cross_entropy_with_logits

fn binary_cross_entropy_with_logits(x, y)expression

Binary cross-entropy of the logits x against the targets y in [0, 1], each row's mean, in the form that neither overflows nor loses a small probability: max(x, 0) - x y + log(1 + e^-|x|).

Source, std/loss.ph:41

fn binary_cross_entropy_with_logits(x, y) = rowmean(max(x, 0.0) - x * y + log(1.0 + exp(-abs(x))))

kl_div

fn kl_div(logp, q)expression

The KL divergence of each row of q from p, p given by its log-probabilities logp: sum q (log q - logp), a q of 0 adding 0. The mean over the rows is PyTorch's reduction="batchmean".

Source, std/loss.ph:36

fn kl_div(logp, q) = rowsum(where(q > 0.0, q * (log(q) - logp), 0.0))

Regression

The regression losses are expression functions, each a row mean of an elementwise term, so the difference and the term fuse into the reduction’s sweep.

mse_loss

fn mse_loss(a, b)expression

The mean squared error of each row: mean((a - b)^2).

Source, std/loss.ph:20

fn mse_loss(a, b) = rowmean(square(a - b))

l1_loss

fn l1_loss(a, b)expression

The mean absolute error of each row: mean(|a - b|).

Source, std/loss.ph:23

fn l1_loss(a, b) = rowmean(abs(a - b))

smooth_l1_loss

fn smooth_l1_loss(a, b, beta)expression

Each row's mean of 0.5 d^2 / beta where |d| < beta and |d| - 0.5 beta elsewhere, d = a - b: the L1 loss with a quadratic bowl around 0, Fast R-CNN's. beta > 0.

Source, std/loss.ph:28

fn smooth_l1_loss(a, b, beta) = rowmean(_smooth_l1(a - b, beta))