python import torch from torch_geometric.data import Data from torch_geometric.nn import GCNConv import numpy as np

模拟数据：100种药物，每种128维特征，已知DDI边200条

num_drugs = 100 features = torch.randn(num_drugs, 128) edge_index = torch.randint(0, num_drugs, (2, 200)) # 无向需添加反向边 edge_index = torch.cat([edge_index, edge_index.flip(0)], dim=1) # 对称化

划分训练/测试边

num_edges = edge_index.shape[1] perm = torch.randperm(num_edges) train_edges = edge_index[:, perm[:300]] # 用作消息传递和损失 test_edges = edge_index[:, perm[300:380]]

负采样：随机生成不存在边

neg_train_edges = torch.randint(0, num_drugs, (2, 300)) neg_test_edges = torch.randint(0, num_drugs, (2, 80))

data = Data(x=features, edge_index=train_edges)

### 定义GNN模型与预测器
```python
class DDIPredictor(torch.nn.Module):
    def __init__(self, in_channels, hidden_channels):
        super().__init__()
        self.conv1 = GCNConv(in_channels, hidden_channels)
        self.conv2 = GCNConv(hidden_channels, hidden_channels)
        self.classifier = torch.nn.Bilinear(hidden_channels, hidden_channels, 1)
    
    def forward(self, x, edge_index, edges_to_predict):
        # 编码所有节点
        x = self.conv1(x, edge_index).relu()
        x = self.conv2(x, edge_index)
        # 提取待预测边的嵌入
        h_i = x[edges_to_predict[0]]
        h_j = x[edges_to_predict[1]]
        scores = self.classifier(h_i, h_j).squeeze()
        return scores

训练与评估

model = DDIPredictor(128, 64)
optimizer = torch.optim.Adam(model.parameters(), lr=0.01)
loss_fn = torch.nn.BCEWithLogitsLoss()

def train():
    model.train()
    optimizer.zero_grad()
    pos_score = model(data.x, data.edge_index, train_edges)
    neg_score = model(data.x, data.edge_index, neg_train_edges)
    loss = loss_fn(pos_score, torch.ones_like(pos_score)) + loss_fn(neg_score, torch.zeros_like(neg_score))
    loss.backward()
    optimizer.step()
    return loss.item()

for epoch in range(200):
    loss = train()
    if epoch % 20 == 0:
        print(f'Epoch {epoch}, Loss: {loss:.4f}')