### PyTorch中结合GNN和BERT的实现代码 在自然语言处理领域，图神经网络（Graph Neural Networks, GNNs）与预训练模型如 BERT 的结合能够有效捕捉文本数据中的结构化关系。这种组合通常用于涉及实体间关系的任务，例如知识图谱补全、问答系统以及语义角色标注。 以下是一个简单的基于 PyTorch 的实现示例，展示如何将 GNN 和 BERT 集成到同一个框架中： #### 数据准备阶段 为了使 GNN 能够利用节点特征，我们可以通过 BERT 提取文本嵌入作为节点表示。假设输入是一组句子及其对应的邻接矩阵，则可以按照如下方式加载数据： ```python import torch from transformers import BertTokenizer, BertModel from torch_geometric.data import Data class TextGraphDataset(torch.utils.data.Dataset): def __init__(self, texts, adjacency_matrix, labels=None): tokenizer = BertTokenizer.from_pretrained('bert-base-uncased') self.text_embeddings = [] for text in texts: inputs = tokenizer(text, return_tensors="pt", padding=True, truncation=True) with torch.no_grad(): bert_model = BertModel.from_pretrained('bert-base-uncased').eval() outputs = bert_model(**inputs).last_hidden_state[:, 0, :] # 取[CLS] token embedding self.text_embeddings.append(outputs.squeeze()) self.adjacency_matrix = torch.tensor(adjacency_matrix, dtype=torch.float32) if labels is not None: self.labels = torch.tensor(labels, dtype=torch.long) def __len__(self): return len(self.text_embeddings) def __getitem__(self, idx): node_features = self.text_embeddings[idx] edge_index = (self.adjacency_matrix > 0).nonzero().t() # 构建边索引 data = Data(x=node_features.unsqueeze(0), edge_index=edge_index.contiguous()) # 创建Data对象 if hasattr(self, 'labels'): label = self.labels[idx] return data, label else: return data ``` #### 模型定义部分 接下来定义一个融合了 GNN 层和线性分类层的模型架构。这里采用 `GCNConv` 来演示基本流程[^1]: ```python import torch.nn.functional as F from torch_geometric.nn import GCNConv class GNNGraphClassifier(torch.nn.Module): def __init__(self, input_dim=768, hidden_channels=128, num_classes=2): super(GNNGraphClassifier, self).__init__() self.conv1 = GCNConv(input_dim, hidden_channels) self.conv2 = GCNConv(hidden_channels, hidden_channels) self.classifier = torch.nn.Linear(hidden_channels, num_classes) def forward(self, data): x, edge_index = data.x, data.edge_index x = F.relu(self.conv1(x, edge_index)) x = F.dropout(x, training=self.training) x = F.relu(self.conv2(x, edge_index)) pooled_x = torch.mean(x, dim=0) # 对所有节点进行平均池化 out = self.classifier(pooled_x) return out ``` #### 训练过程 最后编写训练循环逻辑以便优化参数并评估性能表现： ```python def train(model, optimizer, criterion, dataloader): model.train() total_loss = 0 for batch_idx, (data, target) in enumerate(dataloader): optimizer.zero_grad() output = model(data) loss = criterion(output.view(-1), target) loss.backward() optimizer.step() total_loss += loss.item() avg_loss = total_loss / (batch_idx + 1) return avg_loss def evaluate(model, dataloader): model.eval() correct = 0 total = 0 with torch.no_grad(): for data, target in dataloader: output = model(data) _, predicted = torch.max(output.data, 1) total += target.size(0) correct += (predicted == target).sum().item() accuracy = correct / total * 100. return accuracy ``` 通过以上方法即可完成从数据预处理到最终端到端训练整个工作流的设计开发[^2].