| from __future__ import annotations |
|
|
| import torch |
| from torch import nn |
| from torch.nn import functional as F |
|
|
|
|
| class ContrastiveEncoder(nn.Module): |
| def __init__(self) -> None: |
| super().__init__() |
| self.features = nn.Sequential( |
| nn.Conv2d(1, 12, kernel_size=3, padding=1), |
| nn.GELU(), |
| nn.Conv2d(12, 24, kernel_size=3, padding=1), |
| nn.GELU(), |
| nn.MaxPool2d(2), |
| nn.Flatten(), |
| nn.Linear(24 * 4 * 4, 64), |
| nn.GELU(), |
| ) |
| self.projector = nn.Sequential( |
| nn.Linear(64, 32), |
| nn.GELU(), |
| nn.Linear(32, 16), |
| ) |
|
|
| def encode(self, pixels: torch.Tensor) -> torch.Tensor: |
| return self.features(pixels) |
|
|
| def forward(self, pixels: torch.Tensor) -> torch.Tensor: |
| return self.projector(self.encode(pixels)) |
|
|
|
|
| def parameter_count(model: nn.Module) -> int: |
| return sum(parameter.numel() for parameter in model.parameters()) |
|
|
|
|
| def nt_xent(first: torch.Tensor, second: torch.Tensor, temperature: float) -> torch.Tensor: |
| batch_size = len(first) |
| representations = F.normalize(torch.cat([first, second]), dim=1) |
| similarities = representations @ representations.T / temperature |
| diagonal = torch.eye(2 * batch_size, dtype=torch.bool) |
| similarities = similarities.masked_fill(diagonal, -1e9) |
| positives = torch.cat( |
| [ |
| torch.arange(batch_size, 2 * batch_size), |
| torch.arange(0, batch_size), |
| ] |
| ) |
| return F.cross_entropy(similarities, positives) |
|
|
