mattn/tensai · 上手攻略

• 仓库：mattn/tensai
• 链接：https://github.com/mattn/tensai
• 分类：AI/ML · Go 机器学习框架
• 作者：Tom
• 更新：2026-08-26

这是什么

tensai（天才）是 mattn（Go 社区知名贡献者）用纯 Go 实现的一个微型神经网络框架，核心亮点是全部用 Go 原生 SIMD/AVX2 指令加速矩阵运算，无需 cgo、无需汇编文件、无需外部依赖（开启 wgpu 标签时仅额外引入一个纯 Go 绑定）。它是目前少数可以在纯 Go 生态内完成从训练到模型导出完整流程的 ML 框架。

解决什么问题

• 不想引入 Python / PyTorch 生态，想在 Go 服务中直接训练或推理小型模型
• 需要在 Go 程序中做 int4 / int8 量化推理（比如把 7B 模型塞进 RAM）
• 需要一个干净、可读、无外部依赖的 SIMD 矩阵运算库
• 研究 Go 原生 SIMD 编程（GOEXPERIMENT=simd）

快速安装

环境要求

⚠️ 平台说明： - AMD64 架构 + GOEXPERIMENT=simd → AVX2 加速版（Matmul、ReLU、Sigmoid、Tanh、GELU、LayerNorm、Adam 等全部 8 通道向量化） - 非 AMD64 或未加标志 → 自动使用 portable 回退版本（纯 Go 逻辑，非 SIMD） - Go 1.26 和 1.27 两个版本的 simd API 差异通过构建标签（build tags）兼容

安装命令

# SIMD 加速版（AMD64）
GOEXPERIMENT=simd go install github.com/mattn/tensai@latest

# 或在项目中使用
GOEXPERIMENT=simd go get github.com/mattn/tensai@latest

⚠️ 注意：tensai 没有正式的 go install 入口可执行文件（不是 CLI 工具），通常作为库引入：

import "github.com/mattn/tensai"

// AMD64 SIMD 构建
// go build -tags "simd1_26" .   // Go 1.26
// 或
// go build -tags "simd1_27" .   // Go 1.27

WebGPU 后端（实验性）

如果想用 GPU 加速 Matmul（跨平台 Vulkan/Metal/D3D12）：

GOEXPERIMENT=simd go build -tags wgpu .

⚠️ wgpu 构建标签需要系统安装 wgpu-native 共享库（运行时 dlopen 加载，无需 cgo）。

核心用法

1. 最简示例：MNIST 分类

package main

import (
    "fmt"
    "github.com/mattn/tensai"
    "github.com/mattn/tensai/_example"
)

func main() {
    // 加载 MNIST 数据集（自动下载）
    train, test, err := _example.MNIST()
    if err != nil {
        panic(err)
    }

    // 构建 Sequential 模型：CNN 版
    model := tensai.NewSequential(
        tensai.NewConv2D(1, 8, 3, 3),
        tensai.NewReLU(),
        tensai.NewMaxPool2D(2, 2),
        tensai.NewConv2D(8, 16, 3, 3),
        tensai.NewReLU(),
        tensai.NewMaxPool2D(2, 2),
        tensai.NewFlatten(),
        tensai.NewDense(16*4*4, 10),
        tensai.NewSoftmax(),
    )

    // 编译：优化器 + 损失函数
    model.Compile(tensai.NewAdam(0.001), tensai.NewSoftmaxCrossEntropy())

    // 训练
    model.Fit(train, 10, 32)  // 10 epochs, batch size 32

    // 推理
    logits := model.Predict(test.X)
    fmt.Println("Accuracy:", tensai.Accuracy(logits, test.Y))
}

⚠️ MNIST 数据集自动从网上下载（来源：_example 包内），无网络环境请提前备好数据集文件。

2. 低层 API：自动微分（micrograd 风格）

package main

import (
    "fmt"
    "github.com/mattn/tensai"
)

func main() {
    x := tensai.NewInput(2)
    w1 := tensai.NewParam(tensai.Randn(2, 4))
    b1 := tensai.NewParam(tensai.Zeros(4))
    y := x.MatMul(w1).Add(b1).ReLU()

    // 反向传播
    y.Backward(tensai.Ones(4))

    fmt.Println("Grad w1:", w1.Grad())
}

3. 量化推理（int4 / int8）

// int4 组级量化（7B 模型可从 ~28GB 压缩到 ~3.5GB）
qweights := tensai.QuantizeMatrix4(weights)

// int8 全整数路径（内存带宽 ~31 GB/s，16核）
qweights := tensai.QuantizeMatrix(weights)

⚠️ int4 / int8 量化目前主要用于推理阶段，训练仍建议用 FP32。

4. 模型序列化

// 保存 Sequential 模型参数（JSON 格式）
model.SaveFile("model.json")

// 加载
model2 := tensai.NewSequential(...)
model2.LoadFile("model.json")

5. 导出模型

tensai 支持三种主流格式导出：

// TFLite 导出
model.ExportTFLite("model.tflite")

// ONNX 导出（opset 13，FP32）
model.ExportONNX("model.onnx")

// Safetensors 读写（F16/BF16/F64 → float32）
tensors := tensai.ReadSafetensors("model.safetensors")

⚠️ 导出 TFLite / ONNX 需要对应包支持，ONNX conv 格式为 NCHW（tensai 自身 row layout），导出时无重排序。

6. GGUF 加载（ llama.cpp 格式）

// 从 GGUF 文件加载模型权重
model, err := tensai.LoadGGUF("model.gguf")
// 支持 Q2_K ~ Q6_K、K-quants、MXFP4 等格式

tensai 的核心架构

张量层次

Matrix（2维）→ Tensor（N维，支持 NumPy 式广播）
           ↓
      Param / Input（参与 / 不参与梯度）

SIMD 加速覆盖的算子

⚠️ 所有 SIMD 内核使用 Go 官方 simd/archsimd 实验包，不依赖任何外部汇编或 C 代码。

低分配训练

单步 MLP 训练仅需 ~29 次内存分配，GC 基本不干扰训练循环：

// 复用前向/反向 scratch buffer
model.FitStep(batch)  // 内部复用缓冲区

典型适用场景

1. Go 服务内嵌小型模型推理：不需要 Python 侧渲染，直接在 Go 程序中跑模型
2. 模型量化压缩研究：int4 / int8 量化、GGUF 格式加载、llama.cpp 互操作
3. SIMD 编程学习：用 Go 写 SIMD 内核的最佳参照实现
4. 嵌入式 / 边缘推理：纯 Go 无 C 依赖，适合容器化部署
5. TFLite / ONNX 导出中转：在 Go 中训练，导出到其他推理引擎

坑与注意

⚠️ 重要注意事项：

1. GOEXPERIMENT=simd 是实验特性：Go 官方尚未稳定，升级 Go 版本时可能需要重新检查 simd API 兼容性（已通过 build tags 兼容 1.26 和 1.27）
2. AMD64 专用：ARM64 / ARMv7 等架构只能用 portable 回退版本，性能差距显著
3. 量化仅支持推理：训练流程仍建议 FP32，量化模型主要用于部署推理
4. 无 GPU 训练支持：即使开启 wgpu 标签，也仅 MatMul 在 GPU 运行，训练本身在 CPU
5. 文档较少：纯自述文件 + 代码注释，无正式文档站，遇到问题需要读源码
6. ONNX 导出仅支持 opset 13：不支持更高级 opset 的算子会导出失败
7. Safetensors 读取是懒加载：大模型文件不会一次性读入内存，适合加载超大权重

与同类对比

*开启 wgpu 标签时额外引入 ebitengine/purego

一句话推荐结论

tensai 是目前 Go 生态中最完整的纯 Go 神经网络框架，SIMD 加速、量化推理、GGUF/ONNX/TFLite 全格式支持让它在小模型推理场景下极具竞争力——在 Go 里做 AI 能力集成的首选底层库，但大规模训练仍建议回 Python 生态。