Kaya - KataGo ONNX 模型

本仓库包含围棋游戏（Baduk/Weiqi）的 KataGo 神经网络模型的 ONNX 转换版本。

这些模型为 Kaya 应用提供支持，这是一个基于 Web 的围棋应用程序，具有 AI 驱动的游戏分析和落子建议功能。

模型描述

这些模型从官方 KataGo PyTorch 检查点转换为 ONNX 格式，用于基于 Web 和跨平台的应用程序。

可用模型

每个模型提供三个版本：

• .fp32.onnx - 全精度（FP32）- 推荐用于浏览器/WASM
• .fp16.onnx - 半精度（FP16）- 适用于原生应用（CoreML、CUDA、WebGPU）
• .uint8.onnx - 量化（UINT8）- 约小 4 倍，适用于内存受限的设备

使用方法

使用 ONNX Runtime (Python)

import onnxruntime as ort
import numpy as np

# 加载模型（浏览器/WASM 使用 .fp32.onnx，原生应用使用 .fp16.onnx）
session = ort.InferenceSession("kata1-b28c512nbt-adam-s11165M-d5387M.fp32.onnx")

# 准备输入 (batch_size, channels, height, width)
bin_input = np.random.randn(1, 22, 19, 19).astype(np.float32)
global_input = np.random.randn(1, 19).astype(np.float32)

# 运行推理
outputs = session.run(None, {
    "bin_input": bin_input,
    "global_input": global_input
})

policy, value, miscvalue, moremiscvalue, ownership, scoring, futurepos, seki, scorebelief = outputs

使用 ONNX Runtime Web (JavaScript)

import * as ort from "onnxruntime-web";

// WASM 后端使用 .fp32.onnx，或使用 .uint8.onnx 以获得更小的下载大小
const session = await ort.InferenceSession.create(
  "kata1-b28c512nbt-adam-s11165M-d5387M.fp32.onnx"
);

const binInput = new ort.Tensor(
  "float32",
  new Float32Array(1 * 22 * 19 * 19),
  [1, 22, 19, 19]
);
const globalInput = new ort.Tensor(
  "float32",
  new Float32Array(1 * 19),
  [1, 19]
);

const results = await session.run({
  bin_input: binInput,
  global_input: globalInput,
});

模型输入

模型输出

原始来源

这些模型源自 David J. Wu (lightvector) 的 KataGo 项目。

• 原始仓库: https://github.com/lightvector/KataGo
• 训练数据: https://katagotraining.org/
• 原始检查点: https://media.katagotraining.org/

许可证

原始 KataGo 神经网络权重在 MIT 许可证下发布。

此 ONNX 转换及相关工具也在 MIT 许可证下发布。

引用

如果您使用这些模型，请引用原始 KataGo 论文：

@article{wu2019accelerating,
  title={Accelerating Self-Play Learning in Go},
  author={Wu, David J.},
  journal={arXiv preprint arXiv:1902.10565},
  year={2019}
}

转换详情

• 转换工具: katago-onnx
• ONNX Opset: 17
• FP16 转换: 内部计算使用 FP16，I/O 保持 FP32 以保证兼容性
• UINT8 量化: 使用 QUInt8 权重的动态量化
• 动态轴: 批次大小、棋盘高度/宽度为动态

致谢

特别感谢：

• David J. Wu (lightvector) 创建了 KataGo
• KataGo 训练社区提供训练好的网络