1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
|
package onnx
//go:generate bash scripts/gen.sh
import (
"encoding/binary"
"math"
"os"
"google.golang.org/protobuf/proto"
"go.jknobloc.com/x/onnx/internal/pb"
)
func ExtractInitializer(name string, initializer string) ([]float32, []int, error) {
var data []byte
if d, err := os.ReadFile(name); err != nil {
return nil, nil, err
} else {
data = d
}
model := &pb.ModelProto{}
if err := proto.Unmarshal(data, model); err != nil {
return nil, nil, err
}
g := model.GetGraph()
var result []float32
var shape []int64
ok := false
for _, init := range g.GetInitializer() {
if init.GetName() == initializer {
raw := init.GetRawData()
result = make([]float32, len(raw)/4)
for i := 0; i < len(result); i++ {
bits := binary.LittleEndian.Uint32(raw[i*4 : (i+1)*4])
result[i] = math.Float32frombits(bits)
}
shape = init.GetDims()
ok = true
break
}
}
if !ok {
panic("initializer not found")
}
shapeInt := make([]int, len(shape))
for i, v := range shape {
shapeInt[i] = int(v)
}
return result, shapeInt, nil
}
|
