File size: 1,242 Bytes
3699b0a | 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 | """Compute token-summed codec loss without weighting short microbatches equally."""
from __future__ import annotations
from dataclasses import dataclass
from torch import Tensor
from torch.nn import functional
from .layout import CODEBOOKS, IGNORE, VOCABULARY
@dataclass(frozen=True)
class CodecLoss:
stream_nats: Tensor
stream_counts: Tensor
@property
def numerator(self) -> Tensor:
return self.stream_nats.sum()
@property
def count(self) -> Tensor:
return self.stream_counts.sum()
def codec_loss(logits: Tensor, labels: Tensor) -> CodecLoss:
if logits.shape[:-1] != labels.shape or logits.shape[-2:] != (CODEBOOKS, VOCABULARY):
message = f'codec loss shape mismatch logits={tuple(logits.shape)}, labels={tuple(labels.shape)}'
raise ValueError(message)
valid = labels.ne(IGNORE)
if not bool(valid.any().item()):
message = 'codec loss has zero supervised targets'
raise ValueError(message)
losses = functional.cross_entropy(logits.float().reshape(-1, VOCABULARY), labels.reshape(-1), ignore_index=IGNORE, reduction='none').reshape(labels.shape)
axes = tuple(range(labels.ndim - 1))
return CodecLoss(losses.sum(dim=axes), valid.sum(dim=axes))
|