File size: 1,573 Bytes
02323ee | 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 | 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)
|