| | import torch |
| | import torch.nn as nn |
| | import torch.nn.functional as F |
| |
|
| |
|
| | class ProjectionHead(nn.Module): |
| | def __init__(self, in_dim: int = 512, out_dim: int = 3008, activation: str = "tanh"): |
| | super().__init__() |
| | self.linear = nn.Linear(in_dim, out_dim, bias=False) |
| | act = activation.lower() |
| | if act == "tanh": |
| | self.act = nn.Tanh() |
| | elif act == "relu": |
| | self.act = nn.ReLU() |
| | elif act == "gelu": |
| | self.act = nn.GELU() |
| | elif act == "sigmoid": |
| | self.act = nn.Sigmoid() |
| | else: |
| | raise ValueError(f"Unsupported activation: {activation}") |
| | nn.init.normal_(self.linear.weight, mean=0.0, std=0.02) |
| |
|
| | def forward(self, x: torch.Tensor) -> torch.Tensor: |
| | return self.act(self.linear(x)) |
| |
|
| |
|
| | def binarize_sign(x: torch.Tensor) -> torch.Tensor: |
| | """Return {0,1} bits: 1 if x>0 else 0, dtype int64. |
| | Expects x on any device; returns tensor on same device. |
| | """ |
| | return (x > 0).to(torch.int64) |
| |
|
| |
|
| | def pack_bits_64(bits: torch.Tensor, dim_features: int) -> torch.Tensor: |
| | """Pack {0,1} bits to int64 words of length 64. |
| | |
| | Args: |
| | bits: Tensor (N, D) with 0/1 int64 entries |
| | dim_features: D, must be divisible by 64 |
| | Returns: |
| | Tensor (N, D//64) int64 on same device. |
| | """ |
| | assert dim_features % 64 == 0, "proj_dim must be divisible by 64 for bit-packing" |
| | if bits.dtype != torch.int64: |
| | bits = bits.to(torch.int64) |
| | N = bits.size(0) |
| | words = dim_features // 64 |
| | bits = bits.view(N, words, 64) |
| | shifts = torch.arange(64, device=bits.device, dtype=torch.int64) |
| | packed = (bits << shifts).sum(-1) |
| | return packed.contiguous() |
| |
|
| |
|
| | def _popcount64(x: torch.Tensor) -> torch.Tensor: |
| | """Compute population count for each int64 element using bit hacks. |
| | Returns same shape tensor with counts in int64. |
| | """ |
| | |
| | m1 = torch.tensor(0x5555555555555555, dtype=torch.int64, device=x.device) |
| | m2 = torch.tensor(0x3333333333333333, dtype=torch.int64, device=x.device) |
| | m4 = torch.tensor(0x0F0F0F0F0F0F0F0F, dtype=torch.int64, device=x.device) |
| | x = x - ((x >> 1) & m1) |
| | x = (x & m2) + ((x >> 2) & m2) |
| | x = (x + (x >> 4)) & m4 |
| | x = x + (x >> 8) |
| | x = x + (x >> 16) |
| | x = x + (x >> 32) |
| | return x & torch.tensor(0x7F, dtype=torch.int64, device=x.device) |
| |
|
| |
|
| | def hamming_distance_packed(a_words: torch.Tensor, b_words: torch.Tensor, block: int = 1024) -> torch.Tensor: |
| | """Compute pairwise Hamming distances between two packed code sets. |
| | |
| | Args: |
| | a_words: (Na, W) int64 packed codes |
| | b_words: (Nb, W) int64 packed codes |
| | block: block size for A to limit memory |
| | Returns: |
| | dist: (Na, Nb) int64 distances |
| | """ |
| | assert a_words.dtype == torch.int64 and b_words.dtype == torch.int64 |
| | Na, W = a_words.shape |
| | Nb, Wb = b_words.shape |
| | assert W == Wb |
| | device = a_words.device |
| | out = torch.empty((Na, Nb), dtype=torch.int64, device=device) |
| | for s in range(0, Na, block): |
| | e = min(Na, s + block) |
| | aw = a_words[s:e] |
| | |
| | xor = aw.unsqueeze(1) ^ b_words.unsqueeze(0) |
| | pc = _popcount64(xor) |
| | dist = pc.sum(-1) |
| | out[s:e] = dist |
| | return out |
| |
|
| |
|
