Choose which device "cuda:x" or "cpu" to use

[AhmedBahget2017]

After installing libtensorflow that supports both CPU and GPU. How to select which device to run on ?
Thanks in advance

package main

import (
	"fmt"

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

func describe(i interface{}) {
	fmt.Printf("(%v, %T)\n", i, i)
}

func main() {
	model := tg.LoadModel("/home/yafoz/Desktop/go/src/goTensorflow/tfModel", []string{"serve"}, nil)

	fakeInput, _ := tf.NewTensor([2][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,
	})

	fmt.Println(len(results))
	for i := 0; i < 1; i++ {
		predictions := results[i]
		dummy := predictions.Value()
		describe(dummy)
		foo, _ := dummy.([][]float32)
		fmt.Println("%V", foo[0][1])
		fmt.Println(predictions.Shape())
		fmt.Println("-----------------")
	}
}

[galeone]

Hi, I recommend doing this operation in Python before exporting the saved model.

# model is your model definition 
@tf.function
def inference(input, on_gpu):
  if tf.equal(on_gpu, True):
    with tf.device("/gpu:0"):
      return model(input)
  return model(input)

[galeone]

Closing this since the answer seems sufficient to fix the problem

[chenyanchen]

Hi, I recommend doing this operation in Python before exporting the saved model.

# model is your model definition 
@tf.function
def inference(input, on_gpu):
  if tf.equal(on_gpu, True):
    with tf.device("/gpu:0"):
      return model(input)
  return model(input)

I'm new for TensorFlow.

Is that mean: use CPU or GPU are depend on saved model? Not a optional for server?