Рекомендательные системы на GPU: DLRM и двухбашенные модели

				
					import torch, torch.nn as nn, torch.nn.functional as F

class Tower(nn.Module):
    def __init__(self, emb_dims, mlp_layers):
        super().__init__()
        self.embs = nn.ModuleDict({k: nn.Embedding(num_embeddings=card, embedding_dim=dim)
                                   for k, (card, dim) in emb_dims.items()})
        dims = [sum(dim for _, dim in emb_dims.values())] + mlp_layers
        self.mlp = nn.Sequential(*[nn.Linear(dims[i], dims[i+1]) for i in range(len(dims)-1)])

    def forward(self, x_sparse):
        e = [self.embs[k](x_sparse[k]) for k in self.embs]  # [B, dim]
        h = torch.cat(e, dim=-1)
        return F.normalize(self.mlp(h), dim=-1)

class TwoTower(nn.Module):
    def __init__(self, user_cfg, item_cfg):
        super().__init__()
        self.user = Tower(user_cfg["emb_dims"], user_cfg["mlp"])
        self.item = Tower(item_cfg["emb_dims"], item_cfg["mlp"])

    def forward(self, user_sparse, item_sparse):
        u = self.user(user_sparse)    # [B, D]
        v = self.item(item_sparse)    # [B, D]
        return u, v

def bpr_loss(u, v):
    # in-batch negatives: scores = u @ v.T
    s = u @ v.T                      # [B, B]
    pos = s.diag()
    neg = s - torch.diag_embed(pos)
    loss = -torch.log(torch.sigmoid(pos.unsqueeze(1) - neg)).mean()
    return loss

# Тренинг-цикл (AMP + FSDP/ZeRO по необходимости)