optimizer.go 2.6 KB
Newer Older
1 2
package pserver

D
dongzhihong 已提交
3
// #cgo CFLAGS: -I ../../
4
// #cgo LDFLAGS: ${SRCDIR}/client/c/libpaddle_go_optimizer.a -lstdc++ -lm
D
dongzhihong 已提交
5 6 7
// #include "paddle/optimizer/optimizer.h"
// #include <stdlib.h>
// #include <string.h>
8
import "C"
D
dongzhihong 已提交
9

10 11 12
import (
	"fmt"
	"unsafe"
D
dongzhihong 已提交
13 14

	log "github.com/sirupsen/logrus"
15 16 17 18 19
)

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

type optimizer struct {
D
dongzhihong 已提交
20 21
	opt         *C.struct_paddle_optimizer
	elementType ElementType
D
dongzhihong 已提交
22 23 24 25 26 27 28 29 30 31 32 33 34
}

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]
35 36
}

D
dongzhihong 已提交
37
func newOptimizer(paramWithConfigs ParameterWithConfig, State []byte) *optimizer {
38
	o := &optimizer{}
D
dongzhihong 已提交
39
	o.elementType = paramWithConfigs.Param.ElementType
D
dzhwinter 已提交
40 41
	p := paramWithConfigs.Param
	c := paramWithConfigs.Config
D
dongzhihong 已提交
42
	s := State
43
	paramBufferSize := C.size_t(len(p.Content) / C.sizeof_float)
D
dongzhihong 已提交
44 45
	log.WithFields(log.Fields{
		"ElementType": p.ElementType,
46
		"ParamSize":   paramBufferSize,
D
dongzhihong 已提交
47
		"ConfigSize":  len(c),
D
dongzhihong 已提交
48
		"StateSize":   len(s),
D
dongzhihong 已提交
49
	}).Info("New Optimizer Created with config:")
D
dongzhihong 已提交
50
	var cbuffer unsafe.Pointer
51 52 53
	cbuffer = C.malloc(paramBufferSize)

	C.memcpy(cbuffer, unsafe.Pointer(&p.Content[0]), paramBufferSize)
D
dongzhihong 已提交
54 55 56 57 58
	var cstate unsafe.Pointer
	if len(s) != 0 {
		cstate = unsafe.Pointer(&s[0])
	}

D
dongzhihong 已提交
59
	o.opt = C.paddle_create_optimizer((*C.uchar)(&c[0]), C.int(len(c)),
60
		C.paddle_element_type(p.ElementType), cbuffer, C.int(paramBufferSize), (*C.char)(cstate), C.int(len(s)))
61 62 63
	return o
}

64
func (o *optimizer) GetWeights() []byte {
D
dongzhihong 已提交
65
	var buffer unsafe.Pointer
W
wuyi05 已提交
66 67
	bufferLen := C.paddle_optimizer_get_weights(o.opt, &buffer)
	return cArrayToSlice(buffer, int(bufferLen)*C.sizeof_float)
D
dongzhihong 已提交
68
}
69

D
dongzhihong 已提交
70 71
func (o *optimizer) GetStates() []byte {
	var cbuffer *C.char
72 73
	cbufferLen := C.paddle_optimizer_get_state(o.opt, &cbuffer)
	return cArrayToSlice(unsafe.Pointer(cbuffer), int(cbufferLen))
D
dongzhihong 已提交
74 75
}

D
dongzhihong 已提交
76
func (o *optimizer) UpdateParameter(g Gradient) error {
D
dongzhihong 已提交
77 78
	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)
79 80
	}

D
dongzhihong 已提交
81
	r := C.paddle_update_parameter(o.opt, C.paddle_element_type(g.ElementType), unsafe.Pointer(&g.Content[0]), C.int(len(g.Content))/C.sizeof_float)
82
	if r != 0 {
H
Helin Wang 已提交
83
		return fmt.Errorf("optimizer update returned error code: %d", r)
84 85 86 87 88 89 90
	}
	return nil
}

func (o *optimizer) Cleanup() {
	if unsafe.Pointer(o.opt) != nullPtr {
		C.paddle_release_optimizer(o.opt)
91
		o.opt = (*C.struct_paddle_optimizer)(nullPtr)
92 93
	}
}