forked from ccf-ai-infra/GPUCodeForces
43 lines
1.4 KiB
Python
43 lines
1.4 KiB
Python
import torch
|
||
import torch.nn as nn
|
||
import torch.nn.functional as F
|
||
|
||
BATCH_SIZE = 256
|
||
FEATURE_DIM = 512
|
||
TEMPERATURE = 0.1
|
||
N_NEGATIVES = BATCH_SIZE * 10
|
||
|
||
|
||
class Model(nn.Module):
|
||
|
||
def __init__(self):
|
||
super().__init__()
|
||
self.temperature = TEMPERATURE
|
||
|
||
def forward(self, query: torch.Tensor, positive: torch.Tensor, negatives: torch.Tensor) -> torch.Tensor:
|
||
# InfoNCE Loss实现
|
||
# (B, D) vs (B, D) -> (B,)
|
||
positive_sim = F.cosine_similarity(query, positive, dim=1) / self.temperature
|
||
|
||
# (B, D) @ (D, N_NEG) -> (B, N_NEG)
|
||
negative_sims = torch.matmul(query, negatives.t()) / self.temperature
|
||
|
||
# 拼接: (B, 1) 和 (B, N_NEG) -> (B, 1 + N_NEG)
|
||
logits = torch.cat([positive_sim.unsqueeze(1), negative_sims], dim=1)
|
||
|
||
# 标签总是 0，因为正样本总是在索引 0
|
||
labels = torch.zeros(query.size(0), dtype=torch.long, device=query.device)
|
||
|
||
loss = F.cross_entropy(logits, labels)
|
||
return loss
|
||
|
||
|
||
def get_inputs():
|
||
query = F.normalize(torch.randn(BATCH_SIZE, FEATURE_DIM, dtype=torch.float32), p=2, dim=1)
|
||
positive = F.normalize(torch.randn(BATCH_SIZE, FEATURE_DIM, dtype=torch.float32), p=2, dim=1)
|
||
negatives = F.normalize(torch.randn(N_NEGATIVES, FEATURE_DIM, dtype=torch.float32), p=2, dim=1)
|
||
return [query, positive, negatives]
|
||
|
||
|
||
def get_init_inputs():
|
||
return [] |