numpy より速い？Go の行列演算ライブラリ nune

はじめに

Go にも numpy の様な行列演算ライブラリは幾つかあり、その中でも gonum が有名です。

先日 GitHub を散歩していたら nune という行列演算ライブラリを見つけたので遊んでみる事にしました。

https://github.com/vorduin/nune/

nune とは

nune も他のライブラリ同様に ndarray なのですが、分散処理する事でパフォーマンスを稼いでいる様です。

	Nune	NumPy	PyTorch
Min	3ms	5ms	5ms
Prod	3ms	10ms	5ms
Div	45ms	15ms	10ms
Abs	10ms	25ms	20ms
Sqrt	12ms	40ms	20ms
Tanh	50ms	55ms	65ms

※ Div だけ妙に遅いのが気になる

NumPy は Python とは言え中身はC言語で書かれている事を考えると、結構いいパフォーマンスが出ていると言えるでしょう。

サンプルコード

iris のロジスティック回帰を nune で書いてみました。

package main

import (
	"bufio"
	"fmt"
	"log"
	"math"
	"math/rand"
	"os"

	"github.com/vorduin/nune"
)

func logisticRegression[T nune.Number](X nune.Tensor[T], y nune.Tensor[T], rate float64, ntrains int) nune.Tensor[T] {
	ws := make([]float64, X.Size(1))
	for i := range ws {
		ws[i] = (rand.Float64() - 0.5) * float64(X.Size(1)/2)
	}
	w := nune.From[T](ws)
	for n := 0; n < ntrains; n++ {
		i := 0
		for i < X.Size(0) {
			x := X.Index(i)
			if len(x.Shape()) == 0 {
				i++
				continue
			}
			pred := softmax(w, x)
			perr := y.Ravel()[i] - pred
			scale := T(rate) * perr * pred * (1 - pred)
			dx := x.Clone().Mul(scale)
			for j := 0; j < x.Size(0); j++ {
				w = w.Add(dx)
			}
			i++
		}
	}
	return w
}

func dot[T nune.Number](a, b nune.Tensor[T]) T {
	sa, sb := a.Shape(), b.Shape()
	if len(sa) == 0 || len(sb) == 0 {
		panic("wrong dimention")
	}
	la, lb := sa[0], sb[0]
	if la != lb {
		panic("wrong dimention")
	}

	var sum T
	for i := 0; i < la; i++ {
		sum += a.Ravel()[i] * b.Ravel()[i]
	}
	return sum
}

func softmax[T nune.Number](w, x nune.Tensor[T]) T {
	v := dot(w, x)
	return T(1.0 / (1.0 + math.Exp(float64(-v))))
}

func predict[T nune.Number](w, x nune.Tensor[T]) T {
	return softmax(w, x)
}

func loadData() ([][]float64, []string, error) {
	f, err := os.Open("iris.csv")
	if err != nil {
		log.Fatal(err)
	}
	defer f.Close()

	var resultV [][]float64
	var resultS []string

	scanner := bufio.NewScanner(f)
	// skip header
	scanner.Scan()
	for scanner.Scan() {
		var f1, f2, f3, f4 float64
		var s string
		n, err := fmt.Sscanf(scanner.Text(), "%f,%f,%f,%f,%s", &f1, &f2, &f3, &f4, &s)
		if n != 5 || err != nil {
			continue
		}
		resultV = append(resultV, []float64{f1, f2, f3, f4})
		resultS = append(resultS, s)
	}

	if err = scanner.Err(); err != nil {
		return nil, nil, err
	}
	return resultV, resultS, nil
}

func shuffle[X any, Y any](xx []X, yy []Y) {
	n := len(xx)
	for i := n - 1; i >= 0; i-- {
		j := rand.Intn(i + 1)
		xx[i], xx[j] = xx[j], xx[i]
		yy[i], yy[j] = yy[j], yy[i]
	}
}

func vocab(nn []string) map[string]int {
	m := make(map[string]int)
	for _, n := range nn {
		if _, ok := m[n]; !ok {
			m[n] = len(m)
		}
	}
	return m
}

func onehot[T nune.Number](aa []string, nn map[string]int) nune.Tensor[T] {
	v := nune.Zeros[T](len(aa))
	for i := 0; i < len(aa); i++ {
		f, ok := nn[aa[i]]
		if ok {
			v.Ravel()[i] = T(f)
		}
	}
	return v.Mul(1 / float64(len(nn)))
}

func main() {
	X, Y, err := loadData()
	if err != nil {
		log.Fatal(err)
	}

	ns := vocab(Y)
	m := make([]string, len(ns))
	for k, v := range ns {
		m[v] = k
	}

	shuffle(X, Y)
	n := 100
	xtrain, ytrain, xtest, ytest := X[:n], Y[:n], X[n:], Y[n:]

	Xn := nune.From[float64](xtrain)
	Yn := onehot[float64](ytrain, ns)

	Wn := logisticRegression(Xn, Yn, 0.01, 5000)

	c := 0
	for i, test := range xtest {
		pred := softmax(Wn, nune.From[float64](test))
		p := int(pred*3 + 0.5)
		fmt.Println(m[p], ytest[i])
		if m[p] == ytest[i] {
			c++
		}
	}
	fmt.Printf("Accuracy %f\n", float64(c)/float64(len(xtest)))
}

比較しやすい様に、以前書いた gonum 版の URL も貼っておきます。

https://github.com/mattn/go-gonum-logisticregression-iris/

ベンチマーク

上記サンプルを学習レート 0.01、学習回数 5000 回で比較します。まず nune

real    19.000
system  5.640
user    16.578

次に gonum

real    0.714
system  0.000
user    0.656

gonum は内部で blas を使っているので速いですね。圧勝です。

おわりに

gonum が圧倒的に速い結果を出してしまっただけのベンチマークですが、nune も登場したばかりのライブラリ(ほんの数日前です)なのでこれからどんどんパフォーマンスが良くなっていくのだと思います。

ちなみにビルドを修正する pull-req を送ったら push 権を頂いてしまったので、時間があったらコントリビュートしてみたいと思います。

追記

Python だとこんな感じでしょうか。

from sklearn import metrics
from sklearn.model_selection import train_test_split
from sklearn.linear_model import LogisticRegression
from pandas import DataFrame

from sklearn.datasets import load_iris


iris = load_iris()
x, y = iris.data, iris.target

X = DataFrame(
    x, columns=['Sepal Length', 'Sepal Width', 'Petal Length', 'Petal, Width'])
Y = DataFrame(y, columns=['Species'])
Y['Species'] = Y['Species'].apply(
    lambda a: ['Setosa', 'Veriscolour', 'Virginica'][a])
x_train, x_test, y_train, y_test = train_test_split(x, y, test_size=0.3)

lr = LogisticRegression(C=0.01, max_iter=5000)

lr.fit(x_train, y_train)
pred = lr.predict(x_test)

print(metrics.accuracy_score(y_test, pred))

実行結果 (ランタイム実行部分も含んでいるけど)

0.9111111111111111
real    2.935
system  1.062
user    2.000

README のベンチマークはどうやって計ったんだろう。(また調べる)