File size: 2,537 Bytes
fb0011a | 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 | import torch
import torch.nn as nn
import numpy as np
class FSQ(nn.Module):
def __init__(self, levels: list[int], dim: int = None):
super().__init__()
self.levels = levels
self.dim = dim if dim is not None else len(levels)
self.n_codes = np.prod(levels)
# Register basis as a buffer so it's saved with the model but not a parameter
self.register_buffer('basis', torch.tensor([np.prod(levels[i+1:]) for i in range(len(levels))] + [1], dtype=torch.int64)[:-1])
def bound(self, z, eps=1e-3):
""" Bound z to [-1 + eps, 1 - eps] to avoid dead zones at boundaries """
return (1 - eps) * torch.tanh(z)
def quantize(self, z):
# z: [..., dim]
# 1. Bound
z = self.bound(z)
# 2. Scale and shift to [0, L-1]
quantized = []
indices = []
for i, L in enumerate(self.levels):
# Scale to [0, L-1]
# z_i is in [-1, 1]
val = (z[..., i] + 1) / 2 * (L - 1)
idx = torch.round(val)
# Straight-Through Estimator
# We want gradients to flow through z, but idx is discrete.
# idx_ste = idx + (val - val.detach()) ? No.
# Standard STE: q = z + (round(z) - z).detach()
# Here: idx is the integer index.
# The quantized value in [-1, 1] is:
recon = idx / (L - 1) * 2 - 1
# STE for reconstruction
recon_ste = z[..., i] + (recon - z[..., i]).detach()
quantized.append(recon_ste)
indices.append(idx.long())
quantized = torch.stack(quantized, dim=-1)
indices = torch.stack(indices, dim=-1)
return quantized, indices
def indices_to_codes(self, indices):
# indices: [..., dim]
# Convert to single code index (if needed)
return (indices * self.basis).sum(dim=-1)
def codes_to_indices(self, codes):
# codes: [...]
indices = []
for i in range(len(self.levels)):
indices.append((codes // self.basis[i]) % self.levels[i])
return torch.stack(indices, dim=-1)
def indices_to_values(self, indices):
# indices: [..., dim]
values = []
for i, L in enumerate(self.levels):
# recon = idx / (L - 1) * 2 - 1
val = indices[..., i].float() / (L - 1) * 2 - 1
values.append(val)
return torch.stack(values, dim=-1)
|