how to train model with tfgo? about tfgo HOT 8 OPEN

func (g *Graph) GetDebugString() string {
	tmp := C.ulong(1)
	graphDebugChar := C.TF_GraphDebugString(g.c, (*C.ulong)(&tmp))
	goString := C.GoString(graphDebugChar)

	return goString
}

model, e := tf.LoadSavedModel("model_dir", []string{"serve"}, nil)
if e != nil {
    panic(e.Error())
}

for _, operation := range model.Graph.Operations() {
    fmt.Println(operation.Name())
}

class GolangModel(tf.Module):
    def __init__(self):
        super().__init__()

        bool_input = k.layers.Input(
            shape=(3,),
            name='bool_input',
            dtype='float32',
            batch_size=10
        )

        output = k.layers.Dense(
            1
        )(bool_input)

        self.model = Model(bool_input, output)
        self._global_step = tf.Variable(0, dtype=tf.int32, trainable=False)
        self._optimizer = k.optimizers.Adam()
        self._loss = k.losses.binary_crossentropy

    @tf.function(
        input_signature=[
            tf.TensorSpec(shape=(None, 3), dtype=tf.float32),
            tf.TensorSpec(shape=(None, 1), dtype=tf.float32),
        ]
    )
    def learn(self, data, labels):
        self._global_step.assign_add(1)
        with tf.GradientTape() as tape:
            loss = self._loss(labels, self.model(data))

        gradient = tape.gradient(loss, self.model.trainable_variables)
        self._optimizer.apply_gradients(zip(gradient, self.model.trainable_variables))
        return {"loss": loss}

    @tf.function(input_signature=[tf.TensorSpec(shape=(None, 3), dtype=tf.float32)])
    def predict(self, data):
        prediction = self.model(data)
        return {"prediction": prediction}

gm = GolangModel()

gm.learn(
    tf.zeros([10, 3], dtype=tf.float32),
    tf.zeros([10, 1], dtype=tf.float32),
)
gm.predict(tf.zeros((10, 3), dtype=tf.float32))

tf.saved_model.save(
    gm,
    "/data/models/gm",
    signatures={
        "learn": gm.learn,
        "predict": gm.predict,
    },
)

        gm, e := tf.LoadSavedModel("/data/models/gm", []string{"serve"}, nil)
	if e != nil {
		errorHandler.Error(e)
		return nil
	}

	boolInput, e := tf.NewTensor([][]float32{{0.5, 0.5, 0.5}, {0, 0, 0}})

	result, e := gm.Session.Run(
		map[tf.Output]*tf.Tensor{
			gm.Graph.Operation("predict_data").Output(0): boolInput,
		},
		[]tf.Output{
			gm.Graph.Operation("StatefulPartitionedCall_1").Output(0),
		},
		nil,
	)
	if e != nil {
		errorHandler.Error(e)
		return e
	}

	floatResults, ok := result[0].Value().([][]float32)
	if !ok {
		fmt.Println("No float results")
		return nil
	}

	fmt.Println(floatResults)

	trainData, e := tf.NewTensor([][]float32{
		{0.5, 0.5, 0.5},
		{0.5, 0.5, 0.5},
		{0.5, 0.5, 0.5},
		{0.5, 0.5, 0.5},
		{0.5, 0.5, 0.5},
		{0, 0, 0},
		{0, 0, 0},
		{0, 0, 0},
		{0, 0, 0},
		{0, 0, 0},
	})
	trainLabels, e := tf.NewTensor([][]float32{
		{1},
		{1},
		{1},
		{1},
		{1},
		{0},
		{0},
		{0},
		{0},
		{0},
	})

	for i := 0; i < 1000; i++ {
		_, e := gm.Session.Run(
			map[tf.Output]*tf.Tensor{
				gm.Graph.Operation("learn_data").Output(0):   trainData,
				gm.Graph.Operation("learn_labels").Output(0): trainLabels,
			},
			[]tf.Output{
				gm.Graph.Operation("StatefulPartitionedCall").Output(0),
			},
			nil,
		)
		if e != nil {
			errorHandler.Error(e)
			return e
		}
	}

	boolTest, e := tf.NewTensor([][]float32{{0.5, 0.5, 0.5}, {0, 0, 0}})

	test, e := gm.Session.Run(
		map[tf.Output]*tf.Tensor{
			gm.Graph.Operation("predict_data").Output(0): boolTest,
		},
		[]tf.Output{
			gm.Graph.Operation("StatefulPartitionedCall_1").Output(0),
		},
		nil,
	)
	if e != nil {
		errorHandler.Error(e)
		return e
	}

	testResults, ok := test[0].Value().([][]float32)
	if !ok {
		fmt.Println("No post training float results")
		return nil
	}

	fmt.Println(testResults)

	os.RemoveAll("gm-trained")
	os.MkdirAll("gm-trained/variables", os.ModePerm)
	savedModel, e := ioutil.ReadFile("/data/models/gm/saved_model.pb")
	if e != nil {
		errorHandler.Error(e)
		return e
	}

	e = ioutil.WriteFile("gm-trained/saved_model.pb", savedModel, os.ModePerm)
	if e != nil {
		errorHandler.Error(e)
		return e
	}

	filenameInput, e := tf.NewTensor("gm-trained/variables/variables")
	if e != nil {
		errorHandler.Error(e)
		return e
	}

	_, e = gm.Session.Run(
		map[tf.Output]*tf.Tensor{
			gm.Graph.Operation("saver_filename").Output(0): filenameInput,
		},
		[]tf.Output{
			gm.Graph.Operation("StatefulPartitionedCall_2").Output(0),
		},
		nil,
	)
	if e != nil {
		errorHandler.Error(e)
		return e
	}

	gmTrained, e := tf.LoadSavedModel("gm-trained", []string{"serve"}, nil)
	if e != nil {
		errorHandler.Error(e)
		return nil
	}

	boolTest, e = tf.NewTensor([][]float32{{0.5, 0.5, 0.5}, {0, 0, 0}})

	test, e = gmTrained.Session.Run(
		map[tf.Output]*tf.Tensor{
			gmTrained.Graph.Operation("predict_data").Output(0): boolTest,
		},
		[]tf.Output{
			gmTrained.Graph.Operation("StatefulPartitionedCall_1").Output(0),
		},
		nil,
	)
	if e != nil {
		errorHandler.Error(e)
		return e
	}

	testResults, ok = test[0].Value().([][]float32)
	if !ok {
		fmt.Println("No post training float results")
		return nil
	}

	fmt.Println(testResults)