Gotch: 将PyTorch的强大功能带入Go语言

gotch

Gotch简介

Gotch是一个为Go语言提供PyTorch C++ API (Libtorch)绑定的开源项目。它为Go开发者提供了一个轻量级的包装器,使他们能够利用PyTorch已经优化的C++张量API和动态计算图功能,在Go中进行深度学习开发。

Gotch的目标是将PyTorch强大的深度学习能力引入Go语言生态系统,让Go开发者也能够方便地进行机器学习和人工智能应用的开发。它提供了符合Go语言习惯的API,同时保留了PyTorch的核心功能和性能优势。

Gotch的主要特性

Gotch提供了丰富的功能,主要包括:

全面的PyTorch张量API支持:Gotch封装了2525个PyTorch张量API,让Go开发者可以进行各种张量运算。
完整的PyTorch动态计算图功能:支持自动微分和反向传播,便于构建和训练神经网络模型。
JIT接口:可以运行使用PyTorch Python API训练和保存的模型。
预训练模型支持:能够加载预训练的PyTorch模型并进行推理。
纯Go API:提供纯Go API用于构建和训练神经网络模型,支持CPU和GPU。
最新的计算机视觉模型:包含多种最新的图像处理模型。
NLP语言模型:在单独的Transformer包中提供了基于Gotch构建的Transformer模型,以及纯Go实现的分词器。

Gotch目前处于积极开发阶段,API可能会有破坏性变更。项目欢迎贡献者提交Pull Request、报告问题或参与讨论。当前版本为v0.9.1。

Gotch的安装

Gotch的安装主要分为两个步骤:安装Libtorch和安装Gotch本身。下面分别介绍CPU版本和GPU版本的安装方法。

CPU版本安装

安装Libtorch(如果系统中已有可用的Libtorch可以跳过此步骤):

wget https://github.com/sugarme/gotch/releases/download/v0.9.0/setup-libtorch.sh
chmod +x setup-libtorch.sh
export CUDA_VER=cpu && bash setup-libtorch.sh

安装完成后,需要更新环境变量。在Debian/Ubuntu系统中,在.bashrc文件中添加或更新以下内容:

export GOTCH_LIBTORCH="/usr/local/lib/libtorch"
export LIBRARY_PATH="$LIBRARY_PATH:$GOTCH_LIBTORCH/lib"
export CPATH="$CPATH:$GOTCH_LIBTORCH/lib:$GOTCH_LIBTORCH/include:$GOTCH_LIBTORCH/include/torch/csrc/api/include"
export LD_LIBRARY_PATH="$LD_LIBRARY_PATH:$GOTCH_LIBTORCH/lib"

安装Gotch:

wget https://github.com/sugarme/gotch/releases/download/v0.9.0/setup-gotch.sh
chmod +x setup-gotch.sh
export CUDA_VER=cpu && export GOTCH_VER=v0.9.1 && bash setup-gotch.sh

GPU版本安装

在安装GPU版本之前,请确保您的机器上已经正确安装了CUDA。可以通过以下命令检查CUDA版本:

nvidia-smi

如果需要,请安装NVIDIA驱动、CUDA和CuDNN。

安装Libtorch(如果系统中已有可用的Libtorch可以跳过此步骤):

wget https://github.com/sugarme/gotch/releases/download/v0.9.0/setup-libtorch.sh
chmod +x setup-libtorch.sh
export CUDA_VER=11.8 && bash setup-libtorch.sh

安装完成后,更新环境变量。在Debian/Ubuntu系统中,在.bashrc文件中添加或更新以下内容:

export GOTCH_LIBTORCH="/usr/local/lib/libtorch"
export LIBRARY_PATH="$LIBRARY_PATH:$GOTCH_LIBTORCH/lib"
export CPATH="$CPATH:$GOTCH_LIBTORCH/lib:$GOTCH_LIBTORCH/include:$GOTCH_LIBTORCH/include/torch/csrc/api/include"
LD_LIBRARY_PATH="$LD_LIBRARY_PATH:$GOTCH_LIBTORCH/lib:/usr/lib64-nvidia:/usr/local/cuda-${CUDA_VERSION}/lib64"

安装Gotch:

wget https://github.com/sugarme/gotch/releases/download/v0.9.0/setup-gotch.sh
chmod +x setup-gotch.sh
export CUDA_VER=11.8 && export GOTCH_VER=v0.9.1 && bash setup-gotch.sh

Gotch使用示例

下面通过两个简单的示例来展示Gotch的基本用法。

基本张量操作

这个示例展示了如何创建张量、进行基本运算以及查看张量信息:

import (
    "fmt"

    "github.com/sugarme/gotch"
    "github.com/sugarme/gotch/ts"
)

func basicOps() {
    // 创建一个随机张量
    xs := ts.MustRand([]int64{3, 5, 6}, gotch.Float, gotch.CPU)
    fmt.Printf("%8.3f\n", xs)
    fmt.Printf("%i", xs)

    // 基本张量运算
    ts1 := ts.MustArange(ts.IntScalar(6), gotch.Int64, gotch.CPU).MustView([]int64{2, 3}, true)
    defer ts1.MustDrop()
    ts2 := ts.MustOnes([]int64{3, 4}, gotch.Int64, gotch.CPU)
    defer ts2.MustDrop()

    mul := ts1.MustMatmul(ts2, false)
    defer mul.MustDrop()

    fmt.Printf("ts1:\n%2d", ts1)
    fmt.Printf("ts2:\n%2d", ts2)
    fmt.Printf("mul tensor (ts1 x ts2):\n%2d", mul)

    // 原位操作
    ts3 := ts.MustOnes([]int64{2, 3}, gotch.Float, gotch.CPU)
    fmt.Printf("Before:\n%v", ts3)
    ts3.MustAddScalar_(ts.FloatScalar(2.0))
    fmt.Printf("After (ts3 + 2.0):\n%v", ts3)
}

这个示例展示了如何创建随机张量、进行矩阵乘法运算,以及执行原位加法操作。通过这些基本操作,我们可以看到Gotch提供了简洁而强大的API来处理张量数据。

简化版卷积神经网络

下面的示例展示了如何使用Gotch构建一个简单的卷积神经网络:

import (
    "fmt"

    "github.com/sugarme/gotch"
    "github.com/sugarme/gotch/nn"
    "github.com/sugarme/gotch/ts"
)

type Net struct {
    conv1 *nn.Conv2D
    conv2 *nn.Conv2D
    fc    *nn.Linear
}

func newNet(vs *nn.Path) *Net {
    conv1 := nn.NewConv2D(vs, 1, 16, 2, nn.DefaultConv2DConfig())
    conv2 := nn.NewConv2D(vs, 16, 10, 2, nn.DefaultConv2DConfig())
    fc := nn.NewLinear(vs, 10, 10, nn.DefaultLinearConfig())

    return &Net{
        conv1,
        conv2,
        fc,
    }
}

func (n Net) ForwardT(xs *ts.Tensor, train bool) *ts.Tensor {
    xs = xs.MustView([]int64{-1, 1, 8, 8}, false)

    outC1 := xs.Apply(n.conv1)
    outMP1 := outC1.MaxPool2DDefault(2, true)
    defer outMP1.MustDrop()

    outC2 := outMP1.Apply(n.conv2)
    outMP2 := outC2.MaxPool2DDefault(2, true)
    outView2 := outMP2.MustView([]int64{-1, 10}, true)
    defer outView2.MustDrop()

    outFC := outView2.Apply(n.fc)
    return outFC.MustRelu(true)
}

func main() {
    vs := nn.NewVarStore(gotch.CPU)
    net := newNet(vs.Root())

    xs := ts.MustOnes([]int64{8, 8}, gotch.Float, gotch.CPU)

    logits := net.ForwardT(xs, false)
    fmt.Printf("Logits: %0.3f", logits)
}