optimizer.go 2.2 KB
Newer Older
1 2
package pserver

D
dongzhihong 已提交
3 4
// #cgo pkg-config: protobuf
// #cgo CFLAGS: -I ../../
D
dongzhihong 已提交
5
// //FIXME: ldflags contain "build" path
D
dongzhihong 已提交
6 7 8 9
// #cgo LDFLAGS: ../../build/go/pserver/cclient/libpaddle_go_optimizer.a -lstdc++
// #include "paddle/optimizer/optimizer.h"
// #include <stdlib.h>
// #include <string.h>
10
import "C"
D
dongzhihong 已提交
11

12 13 14 15 16 17 18 19
import (
	"fmt"
	"unsafe"
)

var nullPtr = unsafe.Pointer(uintptr(0))

type optimizer struct {
20
	opt *C.struct_paddle_optimizer
D
dongzhihong 已提交
21 22 23 24 25 26 27 28 29 30 31 32 33 34 35
	// used in GetParam, reconstruct Parameter from optimizer
	ElementType ElementType
}

func cArrayToSlice(p unsafe.Pointer, len int) []byte {
	if p == nullPtr {
		return nil
	}

	// create a Go clice backed by a C array, reference:
	// https://github.com/golang/go/wiki/cgo#turning-c-arrays-into-go-slices
	//
	// Go garbage collector will not interact with this data, need
	// to be freed properly.
	return (*[1 << 30]byte)(p)[:len:len]
36 37
}

D
dzhwinter 已提交
38
func newOptimizer(paramWithConfigs ParameterWithConfig) *optimizer {
39
	o := &optimizer{}
D
dongzhihong 已提交
40
	o.ElementType = paramWithConfigs.Param.ElementType
D
dzhwinter 已提交
41 42
	p := paramWithConfigs.Param
	c := paramWithConfigs.Config
D
dongzhihong 已提交
43
	var cbuffer unsafe.Pointer
44 45 46
	cbuffer_len := int(unsafe.Sizeof(p.Content[0])) * len(p.Content)
	cbuffer = C.malloc(C.size_t(cbuffer_len))
	C.memcpy(cbuffer, unsafe.Pointer(&p.Content[0]), C.size_t(cbuffer_len))
D
dongzhihong 已提交
47 48 49
	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.char)(nullPtr), 0)
50 51 52
	return o
}

53
func (o *optimizer) GetWeights() []byte {
D
dongzhihong 已提交
54
	var buffer unsafe.Pointer
D
dongzhihong 已提交
55
	buffer_len := C.paddle_optimizer_get_weights(o.opt, &buffer)
56
	return cArrayToSlice(buffer, int(buffer_len))
D
dongzhihong 已提交
57
}
58

D
dongzhihong 已提交
59 60
func (o *optimizer) UpdateParameter(g Gradient) error {
	if o.ElementType != g.ElementType {
D
dongzhihong 已提交
61
		return fmt.Errorf("Name: %s, parameter and gradient element type not match, parameter: %v, gradient: %v", g.Name, o.ElementType, g.ElementType)
62 63
	}

D
dongzhihong 已提交
64
	fmt.Println(g)
65
	r := C.paddle_update_parameter(o.opt, C.paddle_element_type(g.ElementType), unsafe.Pointer(&g.Content[0]), C.int(len(g.Content)))
66
	if r != 0 {
H
Helin Wang 已提交
67
		return fmt.Errorf("optimizer update returned error code: %d", r)
68 69 70 71 72 73 74
	}
	return nil
}

func (o *optimizer) Cleanup() {
	if unsafe.Pointer(o.opt) != nullPtr {
		C.paddle_release_optimizer(o.opt)
75
		o.opt = (*C.struct_paddle_optimizer)(nullPtr)
76 77
	}
}