GNN на GPU: GraphSAGE/GAT и выборка

				
					import torch, torch.nn as nn, torch.nn.functional as F
from torch.cuda.amp import autocast, GradScaler
# Предполагаем наличие loader, который выдаёт батчи k-hop подграфов:
# for seeds, blocks, labels in loader: blocks = [ (x, edge_index), ... ]

class SAGE(nn.Module):
    def __init__(self, in_dim, hidden, out_dim, layers=2):
        super().__init__()
        self.lins = nn.ModuleList()
        dims = [in_dim] + [hidden]*(layers-1) + [out_dim]
        for a, b in zip(dims[:-1], dims[1:]):
            self.lins.append(nn.Linear(a, b))
    def forward(self, x, blocks):
        h = x
        for i, (x_src, edge_index) in enumerate(blocks):
            h_dst = h[:x_src.shape[0]]  # convention: dst first
            h = self.lins[i](h)
            h = F.relu(h)
            # mean aggregation (псевдокод)
            # h = aggregate_mean(h, edge_index)
        return h

class GAT(nn.Module):
    def __init__(self, in_dim, hidden, out_dim, layers=2, heads=4):
        super().__init__()
        self.layers = nn.ModuleList()
        for i in range(layers):
            h_in = in_dim if i==0 else hidden*heads
            h_out = out_dim if i==layers-1 else hidden
            self.layers.append(nn.MultiheadAttention(embed_dim=h_out*heads, num_heads=heads, batch_first=True))
    def forward(self, x, blocks):
        h = x
        for attn, (x_src, edge_index) in zip(self.layers, blocks):
            # attention over neighborhood (псевдокод: соберите соседей по edge_index)
            h, _ = attn(h, h, h, need_weights=False)
            h = F.elu(h)
        return h

def train(model, loader, feats, optimizer, epochs=5, device="cuda"):
    scaler = GradScaler()
    model.to(device).train()
    for ep in range(epochs):
        for seeds, blocks, labels in loader:
            x = gather_features(feats, blocks).to(device, non_blocking=True)  # UVA/pinned
            labels = labels.to(device, non_blocking=True)
            optimizer.zero_grad(set_to_none=True)
            with autocast(enabled=True):
                out = model(x, blocks)
                loss = F.binary_cross_entropy_with_logits(out, labels.float())  # для link prediction
            scaler.scale(loss).backward()
            scaler.step(optimizer); scaler.update()