GNN 生产环境部署实战：推理服务构建、图扩展与性能优化

1. 从单机到生产：GNN 应用落地的最后一公里

搞过图神经网络的朋友大概都有这种体会：在 Jupyter Notebook 里用 PyG 或者 DGL 跑通一个 Cora 数据集的节点分类，准确率刷到 80% 以上，感觉一切都很美好。但一旦要把这个模型塞进真实业务里——比如社交网络的好友推荐、电商场景的异常交易识别、或者分子属性预测的批量推理——各种问题就全冒出来了。显存不够、邻居采样太慢、批处理逻辑写错、线上延迟飙到几百毫秒，这些都是我过去两年在多个 GNN 项目里反复踩过的坑。

这个系列的前四篇，我们分别搞定了图数据的表示与加载、经典 GNN 层的原理与实现、训练流程的调优、以及模型评估与可视化。到了这第五篇，主题就一个：怎么把训练好的 GNN 从实验室搬到生产环境，并且让它扛得住真实流量。具体来说，我会围绕三个核心问题展开——构建一个可复用的 GNN 推理服务、扩展到大规模图上的采样与批处理策略、以及部署环节的工程化细节，包括模型导出、服务封装、性能压测和常见故障排查。

这篇文章适合两类人看：一类是已经能跑通 GNN 训练、但不知道怎么上线的算法工程师；另一类是有后端或 MLOps 背景、需要接手 GNN 模型部署的工程同学。文中涉及的所有代码都基于 Python 生态，主要用到 PyTorch Geometric（PyG）和 FastAPI，这些工具在社区里成熟度足够高，遇到问题也容易搜到答案。我不会只贴代码，每个关键选择背后的原因、参数怎么算、坑在哪里，都会讲清楚。

2. 构建推理服务：从模型文件到可用 API

2.1 为什么不能直接把训练脚本改成服务

很多人第一反应是：训练脚本里不是已经有model.eval()和torch.no_grad()了吗？直接包一层 Flask 不就行了？我一开始也是这么想的，结果上线第一天就被打脸。问题出在 GNN 和普通 CNN/RNN 的本质差异上——GNN 的输入不是一个独立的样本，而是一个节点加上它的邻居子图。训练时我们通常把整张图加载到内存里，做全图前向传播；但推理时，如果每来一个请求就把整张大图加载一遍，内存和延迟都受不了。

所以构建推理服务的第一步，是明确推理模式。常见的有三种：

• 全图推理（Full-batch Inference）：一次性对整张图的所有节点做预测，适合离线批处理场景，比如每天凌晨给所有用户重新算一遍 embedding。
• 子图推理（Subgraph Inference）：针对单个节点或一批节点，动态采样其 K 跳邻居，构造子图后前向传播。适合在线实时场景。
• 增量推理（Incremental Inference）：图结构变化时只更新受影响的部分节点，适合动态图场景，实现复杂度最高。

我个人的经验是：在线服务优先选子图推理，离线任务用全图推理。下面重点讲子图推理服务的构建。

2.2 模型导出：别再用 pickle 存整个对象

训练完的模型怎么保存？我见过太多人直接torch.save(model, 'model.pth')，把整个模型对象 pickle 下来。这种做法在部署时是灾难——它把模型类的定义、依赖的第三方库版本全绑死了，换个环境加载就报ModuleNotFoundError。正确做法是只保存 state_dict，模型结构在服务代码里重新定义。

# 训练侧：只保存参数 torch.save(model.state_dict(), 'gnn_model.pt') # 服务侧：重新定义结构后加载 from model import GNNModel # 自己维护的模型定义 model = GNNModel(in_channels=1433, hidden_channels=64, out_channels=7) model.load_state_dict(torch.load('gnn_model.pt', map_location='cpu')) model.eval()

这里有个细节：map_location='cpu'很重要。如果训练在 GPU 上做、服务部署在 CPU 机器上，不加这个参数会直接报错。另外，模型定义文件要作为服务代码的一部分纳入版本管理，不能依赖训练时的临时脚本。

对于更复杂的场景，比如需要跨框架部署（PyTorch 训练、ONNX Runtime 推理），可以用torch.onnx.export导出。但要注意，GNN 的动态邻居采样逻辑很难完整导出到 ONNX，通常只能导出单层的前向计算。我的建议是：如果服务端还是 Python 环境，直接用 PyTorch 加载 state_dict 最省事；只有在对延迟极度敏感、需要 C++ 推理引擎时，才考虑 ONNX 或 TorchScript。

2.3 邻居采样：推理服务的性能瓶颈所在

子图推理的核心是邻居采样。给定一个目标节点，我们要采样它的 K 跳邻居，构造一个子图，然后让 GNN 在这个子图上做前向传播。PyG 提供了NeighborLoader，但它在服务场景下有几个问题：一是每次请求都重新采样，没有缓存；二是采样数量固定，无法根据图的实际稀疏程度自适应。

我自己实现过一个轻量级的采样器，核心逻辑是这样的：

import numpy as np from torch_geometric.data import Data class NeighborSampler: def __init__(self, edge_index, num_neighbors=[10, 10]): self.edge_index = edge_index self.num_neighbors = num_neighbors # 预构建邻接表，避免每次请求都遍历 edge_index self.adj = self._build_adjacency(edge_index) def _build_adjacency(self, edge_index): adj = {} src, dst = edge_index[0].numpy(), edge_index[1].numpy() for s, d in zip(src, dst): adj.setdefault(d, []).append(s) return adj def sample(self, target_nodes): # 逐跳采样，返回子图的节点索引和边索引 nodes = set(target_nodes) for _ in range(len(self.num_neighbors)): new_nodes = set() for n in nodes: neighbors = self.adj.get(n, []) if len(neighbors) > self.num_neighbors[0]: sampled = np.random.choice(neighbors, self.num_neighbors[0], replace=False) else: sampled = neighbors new_nodes.update(sampled) nodes.update(new_nodes) return self._build_subgraph(nodes)

这段代码的关键优化点是预构建邻接表。如果每次请求都从edge_index里筛邻居，时间复杂度是 O(E)，E 是边数，百万级边的图根本扛不住。预构建成字典后，查邻居是 O(1)。代价是内存占用增加，但对于百万节点、千万边的图，邻接表大概占几百 MB，现代服务器完全能承受。

注意：采样数量num_neighbors要和训练时保持一致。如果训练用的是[10, 10]，推理时改成[5, 5]，模型看到的邻居分布变了，预测结果会明显偏移。这个坑我在一个推荐项目里踩过，离线指标正常，线上 CTR 掉了 15%，排查了两天才定位到是采样参数不一致。

2.4 FastAPI 服务封装：异步与批处理

服务框架我选 FastAPI，原因很简单：原生支持异步、自动生成 API 文档、性能比 Flask 高一个量级。下面是一个最小可用的推理服务：

from fastapi import FastAPI from pydantic import BaseModel import torch app = FastAPI() model = load_model() sampler = NeighborSampler(edge_index) class PredictRequest(BaseModel): node_ids: list[int] @app.post("/predict") async def predict(req: PredictRequest): subgraph = sampler.sample(req.node_ids) with torch.no_grad(): logits = model(subgraph.x, subgraph.edge_index) preds = logits.argmax(dim=1).tolist() return {"predictions": preds}

这个版本能跑，但有两个明显问题。第一，没有批处理。如果每秒来 100 个请求，每个请求只查一个节点，GPU 利用率极低。第二，采样是同步阻塞的，采样期间 CPU 在忙，GPU 在等。

改进方案是引入微批处理（Micro-batching）：服务端维护一个请求队列，每 10ms 或队列满 32 个请求时，合并成一个 batch 做一次前向传播。这样 GPU 利用率能提升 5-10 倍。实现上可以用asyncio.Queue加一个后台 worker：

import asyncio queue = asyncio.Queue(maxsize=1000) async def batch_worker(): while True: batch = [] # 等待第一个请求 req = await queue.get() batch.append(req) # 尝试在 10ms 内收集更多请求 try: while len(batch) < 32: req = await asyncio.wait_for(queue.get(), timeout=0.01) batch.append(req) except asyncio.TimeoutError: pass # 合并 batch 做推理 all_nodes = [n for r in batch for n in r.node_ids] subgraph = sampler.sample(all_nodes) with torch.no_grad(): logits = model(subgraph.x, subgraph.edge_index) # 分发结果 offset = 0 for r in batch: n = len(r.node_ids) r.future.set_result(logits[offset:offset+n].argmax(dim=1).tolist()) offset += n

这套微批处理机制是我在一个日请求量千万级的场景里验证过的，P99 延迟从 180ms 降到了 45ms。代价是引入了最多 10ms 的排队延迟，对于绝大多数业务场景完全可以接受。

3. 扩展策略：当图大到单机放不下

3.1 图分区：把大图切成小块

单机内存能装下的图，节点数大概在千万级、边数在亿级。再大怎么办？答案是图分区（Graph Partitioning）。核心思路是把一张大图切成若干子图，每个子图分配到一台机器上，节点如果跨分区，就在多台机器上各存一份副本。

分区算法有很多，工业界常用的是METIS和Random Hash。METIS 追求分区均衡且割边最少，但计算开销大，适合离线一次性分区；Random Hash 就是把节点 ID 对分区数取模，简单粗暴，但割边多，跨分区通信频繁。我的经验是：如果图结构相对稳定、分区数不多（小于 16），用 METIS；如果图动态变化频繁，用 Hash 更实际。

PyG 提供了torch_geometric.distributed模块，支持分布式图存储和采样。但说实话，这套东西配置起来相当复杂，需要自己搭一个 KV 存储（比如 Redis 或 RocksDB）来存节点特征和边索引。如果团队没有专门的分布式系统工程师，我建议先用单机多卡 + 内存映射扛一扛，实在扛不住了再上分布式。

3.2 内存映射：用磁盘换内存

内存映射（Memory Mapping）是个被低估的技巧。原理很简单：把节点特征矩阵存成numpy.memmap文件，操作系统会自动把热点数据缓存在内存里，冷数据留在磁盘上。对于特征维度高、但每次推理只访问少量节点的场景，效果非常好。

import numpy as np # 训练侧：把特征存成 memmap features = np.memmap('node_features.dat', dtype='float32', mode='w+', shape=(num_nodes, feat_dim)) features[:] = node_features_array[:] features.flush() # 服务侧：只读方式打开 features = np.memmap('node_features.dat', dtype='float32', mode='r', shape=(num_nodes, feat_dim)) # 访问方式和普通 numpy 数组一样 node_feat = features[node_id]

我用这招在一个 5000 万节点、128 维特征的图上做过测试，单机 64GB 内存，全量加载需要 25GB，用 memmap 后常驻内存只有 3-4GB，推理延迟增加不到 5ms。代价是首次访问冷数据会有磁盘 IO，所以最好在服务启动后做一次预热，把高频节点的特征提前读一遍。

3.3 采样优化：从 O(E) 到 O(1)

前面提到的邻接表方案，在千万级节点上会占用大量内存。更省内存的做法是CSR（Compressed Sparse Row）格式。CSR 用三个数组表示稀疏矩阵：indptr（每行起始偏移）、indices（列索引）、data（值）。查某个节点的邻居，只需要indices[indptr[n]:indptr[n+1]]，时间复杂度 O(degree)。

from scipy.sparse import csr_matrix # 构建 CSR 格式的邻接矩阵 adj_csr = csr_matrix((np.ones(len(edge_index[0])), (edge_index[0], edge_index[1])), shape=(num_nodes, num_nodes)) # 查节点 n 的邻居 neighbors = adj_csr.indices[adj_csr.indptr[n]:adj_csr.indptr[n+1]]

CSR 的内存占用比字典邻接表小 3-5 倍，而且 scipy 的底层是 C 实现，查询速度更快。唯一的不便是构建时需要一次性拿到所有边，不适合动态增删边的场景。如果图是动态的，可以维护一个 CSR 加一个增量边的缓冲区，定期合并。

4. 部署实战：从容器化到性能压测

4.1 容器化：Docker 镜像怎么瘦身

GNN 服务的 Docker 镜像很容易做到 5GB 以上，因为 PyTorch、PyG、CUDA 运行时加起来就很大。镜像太大带来的问题是：构建慢、推送慢、拉取慢、启动慢。我的瘦身经验有这么几条：

• 用python:3.10-slim作为基础镜像，不要用ubuntu或nvidia/cuda完整版。CUDA 运行时按需安装，只装推理需要的libcudart和libcublas。
• PyTorch 装 CPU 版本或 CUDA 精简版。如果服务跑在 CPU 上，直接pip install torch --index-url https://download.pytorch.org/whl/cpu，能省 2GB。
• 多阶段构建。第一阶段装编译依赖、编译 PyG 的 C++ 扩展；第二阶段只拷贝编译好的.so文件和 Python 包。
• 清理 pip 缓存和 apt 缓存。pip install --no-cache-dir和apt-get clean是基本操作。

一个典型的瘦身后的 Dockerfile 长这样：

FROM python:3.10-slim AS builder RUN pip install --no-cache-dir torch --index-url https://download.pytorch.org/whl/cpu RUN pip install --no-cache-dir torch-geometric fastapi uvicorn FROM python:3.10-slim COPY --from=builder /usr/local/lib/python3.10/site-packages /usr/local/lib/python3.10/site-packages COPY --from=builder /usr/local/bin /usr/local/bin COPY app /app WORKDIR /app CMD ["uvicorn", "main:app", "--host", "0.0.0.0", "--port", "8000"]

这样构建出来的镜像大概 1.2GB，比原始版本小了 70%。

4.2 性能压测：别等上线了才发现扛不住

压测是部署前必须做的一步。我常用的工具是Locust和wrk。Locust 用 Python 写压测脚本，灵活；wrk 性能更强，适合纯 HTTP 接口。

压测时重点看三个指标：QPS、P50/P99 延迟、GPU 利用率。我的经验值是：单卡 T4，子图采样 2 跳、每跳 10 个邻居，GNN 两层 GraphSAGE，QPS 大概在 200-300，P99 延迟 50ms 左右。如果 QPS 上不去，先看 GPU 利用率——如果低于 30%，说明瓶颈在采样或数据加载，不在模型计算。

压测脚本示例：

from locust import HttpUser, task, between import random class GNNUser(HttpUser): wait_time = between(0.01, 0.05) @task def predict(self): node_id = random.randint(0, 100000) self.client.post("/predict", json={"node_ids": [node_id]})

跑压测时记得逐步加压，从 10 并发开始，每 2 分钟加 10，观察延迟曲线。如果延迟突然跳升，说明到了系统容量拐点，这时候的 QPS 就是安全上限，线上限流要设在这个值的 70% 左右。

4.3 监控与告警：上线只是开始

服务上线后，必须有的监控指标包括：请求量、延迟分布、错误率、GPU 显存占用、采样队列长度。我推荐用Prometheus + Grafana这套组合，FastAPI 有现成的prometheus-fastapi-instrumentator库，几行代码就能接入。

from prometheus_fastapi_instrumentator import Instrumentator app = FastAPI() Instrumentator().instrument(app).expose(app)

除了这些通用指标，GNN 服务还要特别关注采样命中率和子图规模分布。如果采样命中率突然下降，可能是图结构变了；如果子图规模异常增大，可能是出现了超级节点（度数极高的节点），会导致单次推理延迟飙升。对于超级节点，我的处理方式是限制采样上限，比如每个节点最多采 50 个邻居，超出的部分随机丢弃。

5. 常见问题与排查技巧实录

5.1 推理结果和训练时不一致

这是最高频的问题。原因通常有三个：采样参数不一致、特征归一化不一致、模型处于 train 模式。排查顺序是：先确认model.eval()被调用了，再对比训练和推理的采样参数，最后检查特征预处理代码是否复用。

我遇到过一个隐蔽的 case：训练时用了BatchNorm，推理时忘了切换 eval 模式，导致 batch 统计量用的是推理时的数据，结果全错。这种问题在 CNN 里很常见，GNN 里同样存在。

5.2 显存溢出（OOM）

GNN 推理的显存占用和子图规模成正比。如果出现 OOM，先看是不是遇到了超级节点。可以在采样器里加一个度数检查：

if len(neighbors) > MAX_DEGREE: neighbors = np.random.choice(neighbors, MAX_DEGREE, replace=False)

另一个原因是 batch size 太大。微批处理的 batch 大小要动态调整，根据当前子图平均规模来定。我的做法是维护一个滑动窗口，记录最近 100 次推理的子图节点数，如果平均值超过阈值，就自动减小 batch size。

5.3 服务启动慢

GNN 服务启动时要加载模型、构建邻接表、预热缓存，这些操作加起来可能要几十秒。如果用的是 K8s，健康检查的超时时间要设得足够长，否则容器还没启动完就被 kill 了。我的配置是initialDelaySeconds: 60、timeoutSeconds: 10、failureThreshold: 6。

另外，邻接表的构建可以离线做，把构建好的 CSR 文件打包进镜像，服务启动时直接加载，能省 80% 的启动时间。

5.4 常见问题速查表

5.5 几个我踩过的坑

第一个坑是用 Redis 存邻接表。一开始觉得 Redis 查询快，就把邻接表塞进去了。结果发现网络往返延迟比本地内存查询高两个数量级，QPS 直接掉到个位数。后来改成进程内缓存，性能才回来。教训是：GNN 采样是高频操作，数据必须放在进程内存里，不能走网络。

第二个坑是忽略了特征的数据类型。训练时特征矩阵是float32，部署时为了省内存改成了float16，结果模型精度掉了 3 个点。GNN 对特征精度比 CNN 更敏感，因为邻居聚合会放大误差。除非显存实在不够，否则坚持用 float32。

第三个坑是没有做请求去重。线上流量里经常有重复请求，比如同一个用户短时间内多次刷新。如果每个请求都重新采样、重新推理，纯属浪费。加一层 LRU 缓存，key 是节点 ID 列表的哈希，value 是预测结果，命中率能到 30% 以上。

6. 写在最后：一些个人体会

GNN 部署这件事，难点从来不在模型本身，而在工程细节。采样逻辑、批处理、内存管理、监控告警，每一项单独看都不复杂，但组合在一起就容易出问题。我的建议是：先用最简单的方案跑通全流程，再逐步优化。不要一上来就搞分布式、搞 ONNX、搞微批处理，先把单机服务跑稳，压测出瓶颈，再针对性地解决。

另外，GNN 和传统深度学习模型最大的区别是它对图结构有强依赖。图变了，模型行为就可能变。所以线上一定要监控图的结构指标，比如平均度数、节点数、边数。这些指标异常时，即使服务本身没报错，预测质量也可能已经下降了。

这个系列到这里就告一段落了。从图数据加载到模型训练，再到今天的部署上线，我们走完了一个 GNN 项目的完整生命周期。如果你跟着做下来，应该已经有一套能跑在生产环境里的 GNN 服务了。后续如果要做动态图、异构图、或者多模态图，思路是类似的，只是在采样和模型结构上要做相应调整。有机会再单独开篇聊。