From 5ef1425adb75eb1b0212518e0f12fefd8d9a8970 Mon Sep 17 00:00:00 2001
From: dongzhihong <dzhwinter@gmail.com>
Date: Mon, 3 Jul 2017 21:13:20 +0800
Subject: [PATCH] "init saving model"

---
 go/pserver/optimizer.go | 16 ++++++++++++++--
 go/pserver/service.go   | 12 +++++++++---
 2 files changed, 23 insertions(+), 5 deletions(-)

diff --git a/go/pserver/optimizer.go b/go/pserver/optimizer.go
index b4a040f46..427251f90 100644
--- a/go/pserver/optimizer.go
+++ b/go/pserver/optimizer.go
@@ -40,17 +40,23 @@ func newOptimizer(paramWithConfigs ParameterWithConfig) *optimizer {
 	o.elementType = paramWithConfigs.Param.ElementType
 	p := paramWithConfigs.Param
 	c := paramWithConfigs.Config
+	s := paramWithConfigs.State
 	log.WithFields(log.Fields{
 		"ElementType": p.ElementType,
 		"ParamSize":   len(p.Content),
 		"ConfigSize":  len(c),
+		"StateSize":   len(s),
 	}).Info("New Optimizer Created with config:")
 	var cbuffer unsafe.Pointer
 	cbuffer = C.malloc(C.size_t(len(p.Content)))
 	C.memcpy(cbuffer, unsafe.Pointer(&p.Content[0]), C.size_t(len(p.Content)))
+	var cstate unsafe.Pointer
+	if len(s) != 0 {
+		cstate = unsafe.Pointer(&s[0])
+	}
+
 	o.opt = C.paddle_create_optimizer((*C.uchar)(&c[0]), C.int(len(c)),
-		C.paddle_element_type(p.ElementType), cbuffer, C.int(len(p.Content)/C.sizeof_float),
-		(*C.char)(nullPtr), 0)
+		C.paddle_element_type(p.ElementType), cbuffer, C.int(len(p.Content)/C.sizeof_float), (*C.char)(cstate), C.int(len(s)))
 	return o
 }
 
@@ -60,6 +66,12 @@ func (o *optimizer) GetWeights() []byte {
 	return cArrayToSlice(buffer, int(buffer_len)*C.sizeof_float)
 }
 
+func (o *optimizer) GetStates() []byte {
+	var cbuffer *C.char
+	cbuffer_len := C.paddle_optimizer_get_state(o.opt, &cbuffer)
+	return cArrayToSlice(unsafe.Pointer(cbuffer), int(cbuffer_len))
+}
+
 func (o *optimizer) UpdateParameter(g Gradient) error {
 	if o.elementType != g.ElementType {
 		return fmt.Errorf("Name: %s, parameter and gradient element type not match, parameter: %v, gradient: %v", g.Name, o.elementType, g.ElementType)
diff --git a/go/pserver/service.go b/go/pserver/service.go
index e15a4e5a5..a5ff86290 100644
--- a/go/pserver/service.go
+++ b/go/pserver/service.go
@@ -38,6 +38,7 @@ type Parameter struct {
 type ParameterWithConfig struct {
 	Param  Parameter
 	Config []byte // parameter configuration in Proto Buffer format
+	State  []byte // parameter training state
 }
 
 // Gradient is the gradient of the parameter.
@@ -58,7 +59,7 @@ func NewService(idx int) (*Service, error) {
 	s := &Service{
 		idx: idx,
 	}
-  s.optMap = make(map[string]*optimizer)
+	s.optMap = make(map[string]*optimizer)
 	s.initialized = make(chan struct{})
 	return s, nil
 }
@@ -143,7 +144,12 @@ func (s *Service) GetParam(name string, parameter *Parameter) error {
 // Save tells the parameter server to save parameters.
 func (s *Service) Save(path string, dummy *int) error {
 	<-s.initialized
-
-	// TODO
+	for opt, ok := range s.optMap {
+		if ok != nil {
+			return fmt.Errorf("parameter optimizerMap error: ", ok)
+		}
+		state := opt.GetStates()
+		weights := opt.GetWeights()
+	}
 	return nil
 }
-- 
GitLab