53 lines
2.3 KiB
Python
53 lines
2.3 KiB
Python
import torch
|
|
|
|
from .compile import torch_compile_lazy
|
|
|
|
|
|
@torch_compile_lazy
|
|
def cross_entropy(
|
|
logits: torch.Tensor, targets: torch.Tensor, mask: torch.Tensor, dtype=torch.float32,
|
|
logits_soft_clip: float | None = None) -> torch.Tensor:
|
|
"""Compute cross entropy between multi-codebook targets and model's logits.
|
|
The cross entropy is computed per codebook to provide codebook-level cross entropy.
|
|
Valid timesteps for each of the codebook are pulled from the mask, where invalid
|
|
timesteps are set to 0.
|
|
|
|
Args:
|
|
logits (torch.Tensor): Model's logits of shape [B, K, T, card].
|
|
targets (torch.Tensor): Target codes, of shape [B, K, T].
|
|
mask (torch.Tensor): Mask for valid target codes, of shape [B, K, T].
|
|
dtype (type): Data type of the output cross entropy.
|
|
logits_soft_clip (float): Clipping value for the logits to avoid numerical instability.
|
|
Recommended value: 30.0.
|
|
Returns:
|
|
ce (torch.Tensor): Cross entropy [B, K, T] with type dtype.
|
|
"""
|
|
output_shape = targets.shape
|
|
assert logits.shape[:-1] == targets.shape
|
|
assert mask.shape == targets.shape
|
|
logits = logits.view(-1, logits.shape[-1])
|
|
targets = targets.reshape(-1)
|
|
mask = mask.reshape(-1)
|
|
|
|
safe_targets = torch.where(
|
|
mask,
|
|
targets,
|
|
torch.zeros(1, device=targets.device, dtype=targets.dtype),
|
|
)
|
|
|
|
# Chunking the conversion to float32 to avoid OOMs.
|
|
ce_chunks = []
|
|
for logits_chunk, targets_chunk in zip(torch.chunk(logits, 4), torch.chunk(safe_targets, 4)):
|
|
logits_chunk = logits_chunk.to(dtype)
|
|
if logits_soft_clip is not None:
|
|
logits_chunk = logits_soft_clip * torch.tanh(logits_chunk / logits_soft_clip)
|
|
log_partition = torch.logsumexp(logits_chunk, dim=-1, keepdim=True)
|
|
|
|
# For some reason, the PyTorch cross entropy is super slow with inputs with large cardinality (e.g. 32000)
|
|
# so we reimplement the cross entropy ourselves...
|
|
ce_chunks.append(log_partition - logits_chunk.gather(-1, targets_chunk[..., None]))
|
|
ce = torch.cat(ce_chunks, dim=0)
|
|
ce = ce[..., 0]
|
|
ce = torch.where(mask, ce, torch.zeros(1, device=ce.device, dtype=ce.dtype))
|
|
return ce.view(output_shape)
|