1 files changed, 68 insertions, 0 deletions
diff --git a/onnx/onnx.go b/onnx/onnx.go
new file mode 100644
index 0000000..85db619
--- /dev/null
+++ b/onnx/onnx.go
@@ -0,0 +1,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
+}