| from __future__ import annotations |
|
|
| import torch |
| from torch import nn |
| from torch.nn import functional as F |
|
|
|
|
| class VectorQuantizedAutoencoder(nn.Module): |
| def __init__( |
| self, |
| codebook_size: int = 32, |
| embedding_dimensions: int = 16, |
| ) -> None: |
| super().__init__() |
| self.codebook_size = codebook_size |
| self.embedding_dimensions = embedding_dimensions |
| self.encoder = nn.Sequential( |
| nn.Conv2d(1, 32, kernel_size=4, stride=2, padding=1), |
| nn.SiLU(), |
| nn.Conv2d(32, embedding_dimensions, kernel_size=3, padding=1), |
| ) |
| self.codebook = nn.Embedding(codebook_size, embedding_dimensions) |
| self.decoder = nn.Sequential( |
| nn.ConvTranspose2d( |
| embedding_dimensions, |
| 32, |
| kernel_size=4, |
| stride=2, |
| padding=1, |
| ), |
| nn.SiLU(), |
| nn.Conv2d(32, 1, kernel_size=3, padding=1), |
| nn.Sigmoid(), |
| ) |
| nn.init.uniform_( |
| self.codebook.weight, |
| -1 / codebook_size, |
| 1 / codebook_size, |
| ) |
|
|
| def encode(self, pixels: torch.Tensor) -> torch.Tensor: |
| return self.encoder(pixels) |
|
|
| def quantize( |
| self, |
| encoded: torch.Tensor, |
| ) -> tuple[torch.Tensor, torch.Tensor]: |
| channels_last = encoded.permute(0, 2, 3, 1).contiguous() |
| flat = channels_last.reshape(-1, self.embedding_dimensions) |
| distances = ( |
| flat.square().sum(1, keepdim=True) |
| + self.codebook.weight.square().sum(1) |
| - 2 * flat @ self.codebook.weight.t() |
| ) |
| indices = distances.argmin(dim=1) |
| quantized = self.codebook(indices).reshape(channels_last.shape) |
| quantized = quantized.permute(0, 3, 1, 2).contiguous() |
| straight_through = encoded + (quantized - encoded).detach() |
| return straight_through, indices.reshape(encoded.shape[0], 4, 4) |
|
|
| def decode(self, latent: torch.Tensor) -> torch.Tensor: |
| return self.decoder(latent) |
|
|
| def decode_indices(self, indices: torch.Tensor) -> torch.Tensor: |
| quantized = self.codebook(indices) |
| quantized = quantized.permute(0, 3, 1, 2).contiguous() |
| return self.decode(quantized) |
|
|
| def forward( |
| self, |
| pixels: torch.Tensor, |
| ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: |
| encoded = self.encode(pixels) |
| quantized_st, indices = self.quantize(encoded) |
| quantized = self.codebook(indices).permute(0, 3, 1, 2).contiguous() |
| reconstruction = self.decode(quantized_st) |
| codebook_loss = F.mse_loss(quantized, encoded.detach()) |
| commitment_loss = F.mse_loss(encoded, quantized.detach()) |
| return reconstruction, indices, codebook_loss, commitment_loss |
|
|
|
|
| class ConditionalCodePrior(nn.Module): |
| def __init__( |
| self, |
| codebook_size: int = 32, |
| token_dimensions: int = 32, |
| hidden_dimensions: int = 64, |
| ) -> None: |
| super().__init__() |
| self.codebook_size = codebook_size |
| self.start_token = codebook_size |
| self.token_embedding = nn.Embedding(codebook_size + 1, token_dimensions) |
| self.label_embedding = nn.Embedding(10, 16) |
| self.position_embedding = nn.Embedding(16, 16) |
| self.recurrent = nn.GRU( |
| token_dimensions + 32, |
| hidden_dimensions, |
| batch_first=True, |
| ) |
| self.output = nn.Linear(hidden_dimensions, codebook_size) |
|
|
| def forward(self, input_tokens: torch.Tensor, labels: torch.Tensor) -> torch.Tensor: |
| batch, length = input_tokens.shape |
| positions = torch.arange(length, device=input_tokens.device) |
| token_features = self.token_embedding(input_tokens) |
| condition = self.label_embedding(labels)[:, None, :].expand(batch, length, -1) |
| position = self.position_embedding(positions)[None, :, :].expand(batch, -1, -1) |
| hidden, _ = self.recurrent( |
| torch.cat([token_features, condition, position], dim=2) |
| ) |
| return self.output(hidden) |
|
|
| @torch.inference_mode() |
| def generate( |
| self, |
| labels: torch.Tensor, |
| *, |
| seed: int, |
| temperature: float = 1.0, |
| ) -> torch.Tensor: |
| generator = torch.Generator(device=labels.device).manual_seed(seed) |
| sequence = torch.full( |
| (len(labels), 1), |
| self.start_token, |
| dtype=torch.long, |
| device=labels.device, |
| ) |
| for _ in range(16): |
| logits = self(sequence, labels)[:, -1] |
| if temperature <= 0.05: |
| token = logits.argmax(dim=1, keepdim=True) |
| else: |
| probabilities = torch.softmax(logits / temperature, dim=1) |
| token = torch.multinomial( |
| probabilities, |
| 1, |
| generator=generator, |
| ) |
| sequence = torch.cat([sequence, token], dim=1) |
| return sequence[:, 1:].reshape(-1, 4, 4) |
|
|
|
|
| class TinyVisionJudge(nn.Module): |
| def __init__(self) -> None: |
| super().__init__() |
| self.features = nn.Sequential( |
| nn.Conv2d(1, 8, kernel_size=3, padding=1), |
| nn.GELU(), |
| nn.Conv2d(8, 8, kernel_size=3, padding=1, groups=8), |
| nn.GELU(), |
| nn.Conv2d(8, 12, kernel_size=1), |
| nn.GELU(), |
| nn.MaxPool2d(2), |
| ) |
| self.classifier = nn.Sequential( |
| nn.Flatten(), |
| nn.Linear(12 * 4 * 4, 10), |
| ) |
|
|
| def forward(self, pixels: torch.Tensor) -> torch.Tensor: |
| return self.classifier(self.features(pixels)) |
|
|
|
|
| def parameter_count(model: nn.Module) -> int: |
| return sum(parameter.numel() for parameter in model.parameters()) |
|
|
|
|
