拓十年匠心定制 · 商业建站与技术教学双线并行 咨询热线:400-886-1026 service@lmnt.cn
ARTICLE DETAIL

资讯详情

深耕网站建设与运营推广的一线实战洞察。

五子棋自博弈推理加速116倍:C++与GPU批处理实战

五子棋自博弈推理加速116倍:C++与GPU批处理实战 1. 从一局五子棋说起为什么要死磕推理速度五子棋这东西规则简单到用一张餐巾纸就能讲明白但真要让程序自己跟自己下、下上几万局来训练一个模型事情就完全不一样了。我最近在做一个自博弈self-play的小项目核心思路很朴素让同一个模型分别执黑执白互相对弈用对局结果去更新策略。听起来没什么门槛可一旦把规模拉起来瓶颈立刻暴露——推理速度。最初的版本我用的是纯 Python 写的推理逻辑模型是一个不算大的卷积网络棋盘 15x15。单局对弈大概要跑 30 到 50 步每一步都要做一次前向推理。实测下来一局大概要 1.2 秒左右跑 1000 局自博弈差不多要 20 分钟。这个速度用来验证想法勉强够但想真正训练出点东西动辄需要几十万局时间成本直接爆炸。于是我开始琢磨能不能把推理这一环从 Python 里拿出来用 C 重写再配合 GPU 做批量化推理折腾了大概两周最终把自博弈对局的吞吐从原来的基准提升到了116 倍。这个数字不是理论峰值是端到端跑完整自博弈循环的实测结果。这篇文章就把整个思路、选型、踩过的坑和关键代码结构完整拆一遍给同样在做棋类 AI、自博弈训练或者任何小模型高频推理场景的朋友一个可复现的参考。先说清楚适合谁看如果你写过一点 C对深度学习推理有基本概念正在被 Python 推理的性能问题折磨那这篇基本就是为你写的。如果你完全没碰过 C也不用急着关掉我会把每个关键决策背后的为什么讲清楚你至少能理解这套方案的取舍逻辑迁移到自己的技术栈里。2. 整体设计为什么是 C 加 GPU 批推理2.1 先定位瓶颈到底在哪动手之前我做了一件事把一局对弈的时间拆开看。结果很反直觉——模型前向推理本身只占了大约 35% 的时间剩下的大头全在 Python 的解释器开销、张量创建销毁、以及 GIL 带来的串行化上。也就是说真正的敌人不是模型太慢而是调度太慢。这个发现直接决定了方案方向。如果瓶颈在模型计算量那应该去换更小的模型或者做量化但瓶颈在调度那就必须换语言、换执行模型。Python 每调用一次推理都要经历一次 Python 层到 C 层的跨界、一次 tensor 的构造、一次内存分配这些固定开销在小模型高频调用的场景下被无限放大。五子棋的模型小、单步计算量低恰恰是这种调度开销主导的典型场景。2.2 三个核心决策基于上面的判断我定了三条主线第一推理核心用 C 重写。C 没有解释器开销没有 GIL内存可以自己管对象可以复用。对于每秒钟要调用几千次的场景这是最直接的收益来源。第二GPU 做批量推理。单步推理时 GPU 其实是被浪费的——一个 15x15 的输入算力根本喂不饱。但如果把多个对局的多步状态攒成一个 batch 一起送进去GPU 的吞吐优势才能真正发挥。这里的关键是批处理策略后面会详细讲。第三自博弈的调度逻辑也放进 C。如果 C 只做推理、Python 还负责对局循环那跨界开销依然存在。索性把整个对局循环、状态管理、胜负判定全部下沉到 CPython 只负责启动和收结果。2.3 为什么不用现成的推理引擎有人会问为什么不直接用 ONNX Runtime 或者 TensorRT 这类成熟引擎我试过 ONNX Runtime它的 C API 确实能用但在这个场景下有两个问题一是它对超小 batch、超高频率的调用模式优化有限每次 session run 仍有固定开销二是我想做的批处理是跨对局动态攒批需要在推理层之上自己控制调度用现成引擎反而多了一层抽象。最后我选择用 LibTorchPyTorch 的 C 前端直接加载 TorchScript 模型控制粒度最细也方便和训练侧的 Python 代码共享模型定义。提示如果你的模型已经导出成 ONNX且对延迟不极端敏感ONNX Runtime 依然是更省事的选择。我选 LibTorch 纯粹是因为需要极致的调度控制。3. 核心细节批处理、内存复用与线程模型3.1 动态攒批把碎请求拼成大 batch这是整个方案里最关键的优化点值得单独讲透。自博弈过程中每一局对弈在任意时刻都只需要一次推理当前局面评估。如果每局单独推理GPU 利用率极低。我的做法是维护一个推理请求队列所有活跃对局把当前需要的推理请求丢进队列一个专门的推理线程从队列里取出请求攒到一定数量或者等待一个极短的时间窗口比如 1 毫秒后拼成一个 batch 送进 GPU。这里有个参数需要权衡batch 大小和等待时间。batch 越大GPU 吞吐越高但每局要等更久才能拿到结果等待时间越长攒的 batch 越大但延迟上升。我实测下来batch 设在 64 到 256 之间、等待窗口 1 毫秒是一个比较舒服的平衡点。batch 再往上单次推理时间增长开始抵消吞吐收益等待窗口再长对局循环的响应性明显变差。具体实现上我用了一个简单的双缓冲结构请求队列和结果队列分离推理线程只负责取请求、拼 batch、算、写结果对局线程只负责发请求、等结果、走下一步。两边通过条件变量同步避免忙等。3.2 内存复用别让分配拖后腿C 重写之后如果每次推理都新建 tensor那和 Python 的毛病一样。所以我在推理层预分配了固定大小的输入输出缓冲区每次攒批时只往已有缓冲区里填数据推理完直接读输出。LibTorch 的torch::from_blob可以让你在不拷贝数据的前提下把一块裸内存包装成 tensor这个 API 是性能关键。棋盘状态本身也做了复用。每个对局维护一个固定大小的int8_t数组表示棋盘落子只改一个位置不需要重建整个状态。评估时把这个数组按 batch 拼进输入缓冲区即可。3.3 线程模型对局线程池加单推理线程线程模型我试过两版。第一版是每个对局一个线程推理也各自调用结果线程切换开销巨大GPU 上下文频繁切换性能反而比 Python 还差。第二版改成对局线程池 单推理线程开 N 个对局线程N 取 CPU 核心数左右所有推理请求汇总到一个推理线程。这样 GPU 只有一个消费者上下文稳定批处理也容易做。实测 N 取 8 到 16 之间效果最好再往上对局线程之间的调度开销开始显现。这个数字和你的 CPU 核心数、对局复杂度都有关建议自己压测确定。4. 实操过程从模型导出到跑通自博弈4.1 模型导出TorchScript 是桥梁训练侧还是用 Python 和 PyTorch模型定义不变。关键是导出成 TorchScript这样 C 侧才能加载。导出时要注意两点一是模型 forward 的输入输出必须是明确的 tensor不要有 Python 特有的动态结构二是把 batch 维度显式写进 forward避免 C 侧调用时形状对不上。import torch model.eval() example_input torch.randn(1, 3, 15, 15) traced torch.jit.trace(model, example_input) traced.save(gomoku_model.pt)用trace还是script取决于模型里有没有控制流。五子棋的卷积网络基本是纯前向trace就够了。如果模型里有 if/for 依赖输入的分支得用script。4.2 C 侧加载与推理封装C 侧用 LibTorch 加载模型核心是把攒批-推理-分发这套逻辑封装成一个类。下面是一个简化后的结构重点看批处理的骨架#include torch/script.h #include vector #include queue #include mutex #include condition_variable struct InferRequest { std::vectorint8_t board; // 15*15 棋盘 int player; // 当前执子方 int request_id; }; struct InferResult { std::vectorfloat policy; // 落子概率 float value; // 局面评估 int request_id; }; class BatchInferencer { public: BatchInferencer(const std::string model_path, int max_batch) : max_batch_(max_batch) { module_ torch::jit::load(model_path); module_.eval(); // 预分配输入缓冲区 input_buffer_ torch::zeros({max_batch, 3, 15, 15}, torch::kFloat32); } void submit(InferRequest req) { std::lock_guardstd::mutex lk(mtx_); req_queue_.push(std::move(req)); cv_.notify_one(); } InferResult wait(int request_id) { std::unique_lockstd::mutex lk(res_mtx_); res_cv_.wait(lk, []{ return result_map_.count(request_id); }); auto r result_map_[request_id]; result_map_.erase(request_id); return r; } private: void infer_loop() { while (running_) { std::vectorInferRequest batch; { std::unique_lockstd::mutex lk(mtx_); cv_.wait_for(lk, std::chrono::milliseconds(1), []{ return !req_queue_.empty(); }); while (!req_queue_.empty() batch.size() max_batch_) { batch.push_back(std::move(req_queue_.front())); req_queue_.pop(); } } if (batch.empty()) continue; run_batch(batch); } } void run_batch(std::vectorInferRequest batch) { int n batch.size(); // 填充输入缓冲区不重新分配 auto in input_buffer_.narrow(0, 0, n); for (int i 0; i n; i) { fill_board(in[i], batch[i]); } std::vectortorch::jit::IValue inputs{in}; auto out module_.forward(inputs).toTuple(); auto policy out-elements()[0].toTensor(); auto value out-elements()[1].toTensor(); // 分发结果 for (int i 0; i n; i) { InferResult r; r.request_id batch[i].request_id; r.value value[i].itemfloat(); // policy 拷贝... std::lock_guardstd::mutex lk(res_mtx_); result_map_[r.request_id] std::move(r); res_cv_.notify_all(); } } torch::jit::Module module_; torch::Tensor input_buffer_; int max_batch_; std::queueInferRequest req_queue_; std::mutex mtx_; std::condition_variable cv_; std::mapint, InferResult result_map_; std::mutex res_mtx_; std::condition_variable res_cv_; bool running_ true; };这段代码里几个点值得强调。input_buffer_是预分配的narrow只是取视图不拷贝fill_board直接往视图里写数据。结果分发用了一个 map 加条件变量每个对局线程按 request_id 等自己的结果。这套结构跑起来推理线程几乎不空转GPU 也一直有活干。4.3 对局循环下沉到 C对局循环本身不复杂初始化棋盘、轮流落子、每步调推理、判胜负。关键是把它写成无 Python 依赖的纯 C。胜负判定用最朴素的四方向扫描五子棋棋盘小扫描开销可以忽略。落子选择用推理输出的 policy 加一点温度采样保证自博弈有探索性。这里有个细节自博弈需要一定的随机性否则同一局会反复走出一样的棋。我在 policy 上做了温度采样温度参数随对局步数衰减前期探索、后期收敛。这个策略是从强化学习的常见实践里借来的实测对训练稳定性有帮助。4.4 编译与依赖编译用 CMake链接 LibTorch。LibTorch 的 C ABI 和编译器版本要匹配我踩过一次坑用 GCC 11 编译的 LibTorch 配 GCC 9 的项目链接时报一堆符号错误。解决办法是统一工具链版本或者直接用 LibTorch 官方推荐的编译器。CUDA 版本也要和 LibTorch 的 CUDA 版本对齐否则运行时报找不到 kernel。cmake_minimum_required(VERSION 3.18) project(gomoku_selfplay) set(CMAKE_CXX_STANDARD 17) find_package(Torch REQUIRED) add_executable(selfplay main.cpp inferencer.cpp game.cpp) target_link_libraries(selfplay ${TORCH_LIBRARIES}) set_property(TARGET selfplay PROPERTY CXX_STANDARD 17)编译命令里记得把 LibTorch 的路径传给 CMakecmake -DCMAKE_PREFIX_PATH/path/to/libtorch .. make -j85. 性能实测116 倍是怎么来的5.1 基准与测试环境基准是纯 Python 版本单局对弈约 1.2 秒1000 局约 20 分钟。测试环境是一台带独立显卡的笔记本CPU 是常见的移动端多核GPU 是 RTX 4060 Laptop。C 版本跑同样的 1000 局端到端耗时约 10.3 秒。20 分钟对 10 秒差不多就是 116 倍。这个数字要拆开看不能简单归功于某一项优化。我做了几组对照方案1000 局耗时相对基准纯 Python 单步推理约 1200 秒1xC 单步推理无批处理约 95 秒约 12.6xC 批处理推理batch128约 10.3 秒约 116x可以看到光是把推理搬到 C就拿到了 12 倍左右的提升这部分主要来自消除解释器开销和 GIL。批处理又在此基础上翻了将近 10 倍这部分来自 GPU 利用率的提升。两者叠加才有了 116 倍。5.2 批处理为什么这么猛单步推理时GPU 的算力单元大部分时间在等数据利用率可能只有个位数。batch 到 128 之后同样的 kernel 启动开销被摊薄到 128 个样本上GPU 的并行度被真正喂饱。五子棋模型小单样本计算量低所以批处理的边际收益特别明显。如果换成大模型单样本就能吃满 GPU批处理的收益会小很多——这也是为什么这个方案特别适合小模型高频推理场景。5.3 延迟与吞吐的取舍批处理提升的是吞吐不是单次延迟。对自博弈来说这没问题因为我们要的是单位时间跑更多局单局慢一点无所谓。但如果你的场景是实时对弈人机对战那批处理就不合适了得走单步低延迟路线。这一点在选型时要想清楚别把吞吐优化套到延迟敏感的场景上。6. 常见问题与排查实录6.1 推理结果和 Python 对不上这是最常见的问题。原因通常有三个一是输入数据的归一化方式不一致Python 侧可能做了某种预处理C 侧漏了二是棋盘编码的通道顺序反了比如 Python 是 [黑, 白, 空]C 写成了 [空, 黑, 白]三是浮点精度差异这个一般影响很小但如果模型对数值敏感也可能导致 policy 分布明显不同。排查方法很简单固定一个棋盘状态两边分别推理逐元素对比输出。6.2 GPU 利用率上不去如果发现 GPU 利用率还是很低先检查 batch 是不是真的攒起来了。常见原因是等待窗口太短请求还没攒够就触发了推理。可以把等待窗口调大一点试试或者打印每次实际 batch 的大小。另一个原因是推理线程和对局线程的锁竞争太激烈导致请求入队慢。这种情况可以把请求队列改成无锁队列或者减少锁的粒度。6.3 内存持续增长如果跑久了内存一直涨八成是结果 map 里的条目没被及时清理。对局线程拿到结果后要确保 erase否则 map 会越积越大。另外 LibTorch 的 tensor 如果被意外持有引用也不会释放。用torch::NoGradGuard关掉梯度记录能省不少内存推理场景本来就不需要梯度。6.4 编译链接报错速查报错现象常见原因解决方向大量 undefined symbol编译器 ABI 不匹配统一 GCC 版本找不到 CUDA kernelLibTorch 与 CUDA 版本不一致对齐 CUDA 版本运行时崩溃在 forward输入形状或类型不对检查 tensor 的 shape 和 dtype链接找不到 TorchCMake 路径没配对检查 CMAKE_PREFIX_PATH提示LibTorch 的版本最好和训练侧 PyTorch 版本一致否则 TorchScript 模型可能加载失败。这个坑我踩过模型在 Python 里好好的C 一加载就报版本不兼容。7. 几个让我少走弯路的实操心得第一先做 profiling 再动手。我一开始差点直接去优化模型结构幸好先测了一下发现瓶颈根本不在模型。如果方向错了后面所有努力都是白费。用perf或者简单的计时打点先搞清楚时间花在哪。第二批处理不是越大越好。我试过 batch512结果单次推理时间涨得比吞吐收益还快整体反而变慢。batch 大小要和模型规模、GPU 显存匹配找到那个拐点。我的经验是从 64 开始往上试每次翻倍看端到端吞吐什么时候不再涨。第三别忽视数据搬运。棋盘数据从对局线程搬到推理线程如果每次都深拷贝开销也不小。能用指针或引用就用实在要拷贝也尽量用连续内存加 memcpy别用逐个元素的循环。第四自博弈的随机性要控制好。完全贪心会让对局高度重复训练效果差随机性太大又学不到东西。温度采样加衰减是个比较稳的方案前期多探索后期多利用。第五测试要固定随机种子。自博弈涉及大量随机采样不固定种子的话两次跑的结果没法对比调优会变成玄学。C 侧用std::mt19937配合固定种子Python 侧也同步固定这样两边结果才可比。这套方案跑通之后我把自博弈的规模从几千局提到了几十万局训练迭代速度完全不是一个量级。如果你也在做类似的事情建议先把推理这一环从 Python 里剥出来哪怕先不做批处理光是 C 重写就能拿到十倍左右的提升性价比很高。批处理是第二步等单步推理稳定了再上调试起来也更容易定位问题。
返回列表