Skip to content

Instantly share code, notes, and snippets.

@zhongdai
Last active April 30, 2024 09:41
Show Gist options
  • Save zhongdai/1cc6e0e5c4dc7602908a550d664ea049 to your computer and use it in GitHub Desktop.
Save zhongdai/1cc6e0e5c4dc7602908a550d664ea049 to your computer and use it in GitHub Desktop.
How to install tensorflow go binding on M1

Background

This is the follow up of how to install tensorflow C api on M1. Since I have installed the C API, I need to install the tensorflow Go to run the graph in my Go code. However, I found this document, however, it does not work on a M1 macbook.

The goal of this document is to compile and run the following Go program

package main

import (
	tf "github.com/tensorflow/tensorflow/tensorflow/go"
	"github.com/tensorflow/tensorflow/tensorflow/go/op"
	"fmt"
)

func main() {
	// Construct a graph with an operation that produces a string constant.
	s := op.NewScope()
	c := op.Const(s, "Hello from TensorFlow version " + tf.Version())
	graph, err := s.Finalize()
	if err != nil {
		panic(err)
	}

	// Execute the graph in a session.
	sess, err := tf.NewSession(graph, nil)
	if err != nil {
		panic(err)
	}
	output, err := sess.Run(nil, []tf.Output{c}, nil)
	if err != nil {
		panic(err)
	}
	fmt.Println(output[0].Value())
}

Prerequisites

Compile the C API

Follow this document to ensure you have installed and successful compile/run the sample C program.

Tools to be installed

The Mac OS sed is not good, so install the GNU version gsed,

brew install gsed

The protobuf tools,

brew install protobuf

Steps

Figure out the GOPATH

go env GOPATH

The output is different for each machine, it should look like /Users/Tom/go. I will use /Users/Tom/go in this document, if you have different path, please substitute it in all commands.

Clone the library

git clone --branch v2.17.0 https://github.com/tensorflow/tensorflow.git /Users/Tom/go/src/github.com/tensorflow/tensorflow

Since we have installed 2.17.0 for the C API, we use the same version here.

Change the directorey to /go/src/github.com/tensorflow/tensorflow,

cd /go/src/github.com/tensorflow/tensorflow

Initialse a new go.mod file,

go mod init github.com/tensorflow/tensorflow`

Patch Tensor Standard Library (TSL) protos to declare Go package

noticed we are using gsed in the following steps

gsed -i '4 i option go_package = "github.com\/tensorflow\/tensorflow\/tensorflow\/go\/core\/framework\/dataset_go_proto";' tensorflow/core/framework/dataset.proto
gsed -i '9 c option go_package = "github.com\/tensorflow\/tensorflow\/tensorflow\/go\/core\/framework\/graph_debug_info_go_proto";' tensorflow/core/framework/graph_debug_info.proto
gsed -i '4 i option go_package = "github.com\/tensorflow\/tensorflow\/tensorflow\/go\/core\/framework\/optimized_function_graph_go_proto";' tensorflow/core/framework/optimized_function_graph.proto

Patch Tensor Standard Library (TSL) protos to declare Go package

gsed -i '5 c option go_package = "github.com\/google\/tsl\/tsl\/go_proto";' third_party/xla/third_party/tsl/tsl/protobuf/bfc_memory_map.proto
gsed -i '5 c option go_package = "github.com\/google\/tsl\/tsl\/go_proto";' third_party/xla/third_party/tsl/tsl/protobuf/coordination_config.proto
gsed -i '7 c option go_package = "github.com\/google\/tsl\/tsl\/go_proto";' third_party/xla/third_party/tsl/tsl/protobuf/coordination_service.proto
gsed -i '6 c option go_package = "github.com\/google\/tsl\/tsl\/go_proto";' third_party/xla/third_party/tsl/tsl/protobuf/distributed_runtime_payloads.proto
gsed -i '8 c option go_package = "github.com\/google\/tsl\/tsl\/go_proto";' third_party/xla/third_party/tsl/tsl/protobuf/dnn.proto
gsed -i '12 c option go_package = "github.com\/google\/tsl\/tsl\/go_proto";' third_party/xla/third_party/tsl/tsl/protobuf/error_codes.proto
gsed -i '8 c option go_package = "github.com\/google\/tsl\/tsl\/go_proto";' third_party/xla/third_party/tsl/tsl/protobuf/histogram.proto
gsed -i '5 c option go_package = "github.com\/google\/tsl\/tsl\/go_proto";' third_party/xla/third_party/tsl/tsl/protobuf/rpc_options.proto
gsed -i '10 c option go_package = "github.com\/google\/tsl\/tsl\/go_proto";' third_party/xla/third_party/tsl/tsl/protobuf/status.proto
gsed -i '13 i option go_package = "github.com\/google\/tsl\/tsl\/go_proto";' third_party/xla/third_party/tsl/tsl/protobuf/test_log.proto

Patch tensorflow/go/genop to generate TF and TSL protobufs

gsed -i '71d;72d' tensorflow/go/genop/generate.sh
gsed -i '71 i \    ${TF_DIR}\/third_party\/xla\/xla\/autotuning.proto \\' tensorflow/go/genop/generate.sh
gsed -i '72 i \    ${TF_DIR}\/third_party\/xla\/third_party\/tsl\/tsl\/protobuf\/*.proto; do \\' tensorflow/go/genop/generate.sh
gsed -i '74 i \    -I ${TF_DIR}/third_party/xla/third_party/tsl \\' tensorflow/go/genop/generate.sh
gsed -i '75 i \    -I ${TF_DIR}/third_party/xla \\' tensorflow/go/genop/generate.sh

Figout the libtensorflow Cflags and Lflags

pkg-config --libs --cflags tensorflow

You should see the output like following,

-I/usr/local/lib/libtensorflow-cpu-darwin-arm64-2.15.0/include/tensorflow -L/usr/local/lib/libtensorflow-cpu-darwin-arm64-2.15.0/lib -ltensorflow -ltensorflow_framework

Patch the cgo flags

gsed -i '19 c // #cgo LDFLAGS: -ltensorflow -L/usr/local/lib/libtensorflow-cpu-darwin-arm64-2.15.0/lib' tensorflow/go/lib.go
gsed -i '20 c // #cgo CFLAGS: -I/usr/local/lib/libtensorflow-cpu-darwin-arm64-2.15.0/include/tensorflow' tensorflow/go/lib.go
gsed -i '20 c // #cgo LDFLAGS: -ltensorflow -L/usr/local/lib/libtensorflow-cpu-darwin-arm64-2.15.0/lib' tensorflow/go/genop/internal/lib.go
gsed -i '21 c // #cgo CFLAGS: -I/usr/local/lib/libtensorflow-cpu-darwin-arm64-2.15.0/include/tensorflow' tensorflow/go/genop/internal/lib.go

the content might be different, depends on your output of pkg-config command. Basically, the -I flag indicats the header files, and the -L flag indicates the libtensorflow.* file.

Generate wrappers and protocol buffers

(cd tensorflow/go/op && DYLD_LIBRARY_PATH=/usr/local/lib/libtensorflow-cpu-darwin-arm64-2.15.0/lib go generate)

you should not leave the current directory after this command

Use Go Mod's replace directive to locate TSL protos

go mod edit -require github.com/google/tsl@v0.0.0+incompatible
go mod edit -replace github.com/google/tsl=/Users/Tom/go/src/github.com/google/tsl

Initialize a new go.mod for TSL and add dependencies

(cd /go/src/github.com/google/tsl && go mod init github.com/google/tsl && go mod tidy)

Test the installation

go mod tidy
DYLD_LIBRARY_PATH=/usr/local/lib/libtensorflow-cpu-darwin-arm64-2.15.0/lib go test ./...

Compile and run the Go example

Create the Go Program

mkdir -p /tmp/hello-tf-from-go

Go to this directory, cd /tmp/hello-tf-from-go, then create a file contains the Go program in the above, named it as hello_tf.go

Init the Go mod

$ go mod init hello-tf-from-go
$ go mod edit -require github.com/google/tsl@v0.0.0+incompatible
$ go mod edit -require github.com/tensorflow/tensorflow@v2.15.0+incompatible
$ go mod edit -replace github.com/google/tsl=/Users/Tom/go/src/github.com/google/tsl
$ go mod edit -replace github.com/tensorflow/tensorflow=/Users/Tom/go/src/github.com/tensorflow/tensorflow
$ go mod tidy

Compile and run

go run hello_tf_from_go.go

The output looks like,

2024-04-29 12:53:16.506974: I tensorflow/compiler/mlir/mlir_graph_optimization_pass.cc:388] MLIR V1 optimization pass is not enabled
Hello from TensorFlow version 2.17.0

Reference and Credits

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment