tfgo: TensorFlow in Go

tfgo 是一个简化 TensorFlow 官方 Go 绑定使用的插件库。TensorFlow 的 Go 绑定使用起来较为困难，而 tfgo 使其变得简单！

主要优势

• 解决作用域问题：每个新节点都会有新的唯一名称
• 自动类型转换：属性会自动转换为支持的类型，而不是在运行时抛出错误
• 使用方法链式调用，可以编写更优雅的 Go 代码

依赖

1. TensorFlow-2.9.1 库
2. TensorFlow 绑定 github.com/galeone/tensorflow

安装

go get github.com/galeone/tfgo

快速开始

基础示例

package main

import (
    "fmt"
    tg "github.com/galeone/tfgo"
    tf "github.com/galeone/tensorflow/tensorflow/go"
)

func main() {
    root := tg.NewRoot()
    A := tg.NewTensor(root, tg.Const(root, [2][2]int32{{1, 2}, {-1, -2}}))
    x := tg.NewTensor(root, tg.Const(root, [2][1]int64{{10}, {100}}))
    b := tg.NewTensor(root, tg.Const(root, [2][1]int32{{-10}, {10}}))
    Y := A.MatMul(x.Output).Add(b.Output)
    
    // 注意：Y 只是指向 A 的指针！
    
    // 如果想在图中创建不同的节点，需要克隆 Y 或 A
    Z := A.Clone()
    results := tg.Exec(root, []tf.Output{Y.Output, Z.Output}, nil, &tf.SessionOptions{})
    fmt.Println("Y: ", results[0].Value(), "Z: ", results[1].Value())
    fmt.Println("Y == A", Y == A) // => true
    fmt.Println("Z == A", Z == A) // => false
}

输出：

Y:  [[200] [-200]] Z:  [[200] [-200]]
Y == A true
Z == A false

计算机视觉示例

package main

import (
    tg "github.com/galeone/tfgo"
    "github.com/galeone/tfgo/image"
    "github.com/galeone/tfgo/image/filter"
    "github.com/galeone/tfgo/image/padding"
    tf "github.com/galeone/tensorflow/tensorflow/go"
    "os"
)

func main() {
    root := tg.NewRoot()
    grayImg := image.Read(root, "/home/pgaleone/airplane.png", 1)
    grayImg = grayImg.Scale(0, 255)

    // 使用 Sobel 滤波器进行边缘检测：卷积
    Gx := grayImg.Clone().Convolve(filter.SobelX(root), image.Stride{X: 1, Y: 1}, padding.SAME)
    Gy := grayImg.Clone().Convolve(filter.SobelY(root), image.Stride{X: 1, Y: 1}, padding.SAME)
    convoluteEdges := image.NewImage(root.SubScope("edge"), Gx.Square().Add(Gy.Square().Value()).Sqrt().Value()).EncodeJPEG()

    Gx = grayImg.Clone().Correlate(filter.SobelX(root), image.Stride{X: 1, Y: 1}, padding.SAME)
    Gy = grayImg.Clone().Correlate(filter.SobelY(root), image.Stride{X: 1, Y: 1}, padding.SAME)
    correlateEdges := image.NewImage(root.SubScope("edge"), Gx.Square().Add(Gy.Square().Value()).Sqrt().Value()).EncodeJPEG()

    results := tg.Exec(root, []tf.Output{convoluteEdges, correlateEdges}, nil, &tf.SessionOptions{})

    file, _ := os.Create("convolved.png")
    file.WriteString(results[0].Value().(string))
    file.Close()

    file, _ = os.Create("correlated.png")
    file.WriteString(results[1].Value().(string))
    file.Close()
}

训练与部署示例

Python 训练代码

import tensorflow as tf

model = tf.keras.Sequential(
    [
        tf.keras.layers.Conv2D(
            8,
            (3, 3),
            strides=(2, 2),
            padding="valid",
            input_shape=(28, 28, 1),
            activation=tf.nn.relu,
            name="inputs",
        ),  # 14x14x8
        tf.keras.layers.Conv2D(
            16, (3, 3), strides=(2, 2), padding="valid", activation=tf.nn.relu
        ),  # 7x716
        tf.keras.layers.Flatten(),
        tf.keras.layers.Dense(10, name="logits"),  # linear
    ]
)

tf.saved_model.save(model, "output/keras")

Go 部署代码

package main

import (
    "fmt"
    tg "github.com/galeone/tfgo"
    tf "github.com/galeone/tensorflow/tensorflow/go"
)

func main() {
    model := tg.LoadModel("test_models/output/keras", []string{"serve"}, nil)

    fakeInput, _ := tf.NewTensor([1][28][28][1]float32{})
    results := model.Exec([]tf.Output{
        model.Op("StatefulPartitionedCall", 0),
    }, map[tf.Output]*tf.Tensor{
        model.Op("serving_default_inputs_input", 0): fakeInput,
    })

    predictions := results[0]
    fmt.Println(predictions.Value())
}

TensorFlow 安装

手动安装

curl -L "https://storage.googleapis.com/tensorflow/libtensorflow/libtensorflow-cpu-linux-x86_64-2.9.1.tar.gz" | sudo tar -C /usr/local -xz
sudo ldconfig

Docker 安装

docker pull tensorflow/tensorflow:2.9.1

tfgo 为 Go 开发者提供了更简单的方式来使用 TensorFlow，特别是在需要将训练好的模型部署到生产环境时。通过方法链式调用和自动类型转换等特性，它大大简化了 TensorFlow Go API 的使用。