claude
2c73493bf8
The gate loaded and worked on workpc and took mavwaked from 68% of one core to 335%. onnxruntime sizes its intra-op pool to every core and spins between runs, which an always-on gate scoring three graphs twelve times a second provokes for the whole day. One thread per session brings it to 81%, so the keyword costs about 13% of a core, and each graph still finishes well inside its 80ms. The unit now passes the three -wake- flags and the models sit beside silero_vad.onnx in ~/.local/share/maven/models. The threshold is left at the binary's default so there is one place to change it.
203 lines
6.4 KiB
Go
203 lines
6.4 KiB
Go
package main
|
|
|
|
// The three models behind the wake word (V-487 stage two).
|
|
//
|
|
// openWakeWord's pipeline, run in a row:
|
|
//
|
|
// audio -> melspectrogram.onnx -> 32-bin mel frames, one per 10ms
|
|
// 76 frames -> embedding_model.onnx -> one 96-dim embedding per 80ms
|
|
// 16 embeds -> maven_wakeword.onnx -> one score
|
|
//
|
|
// The first two are frozen and pretrained. Only the last was trained here,
|
|
// which is why it is 100KB and the other two are megabytes. The shapes are
|
|
// not guesses: 2.0s of 16kHz audio measures 197 mel frames, and 76-frame
|
|
// windows at stride 8 give exactly the 16 embeddings the head was fitted on.
|
|
//
|
|
// This file knows ONNX and nothing about the 80ms cadence. wakeword.go knows
|
|
// the cadence and nothing about tensors.
|
|
|
|
import (
|
|
"fmt"
|
|
|
|
ort "github.com/yalue/onnxruntime_go"
|
|
)
|
|
|
|
const (
|
|
// melHop — samples per mel frame. 10ms at 16kHz.
|
|
melHop = 160
|
|
// melBins — mel bins per frame, fixed by melspectrogram.onnx.
|
|
melBins = 32
|
|
// embedFrames — mel frames one embedding is computed over, 760ms.
|
|
embedFrames = 76
|
|
// embedStride — mel frames between embeddings, 80ms.
|
|
embedStride = 8
|
|
// embedDim — the embedding width.
|
|
embedDim = 96
|
|
// headWindow — embeddings the head scores at once, 1.28s of audio.
|
|
headWindow = 16
|
|
|
|
// melContext — samples of history prepended to each incremental mel
|
|
// call, chosen so the eight frames this call yields continue exactly
|
|
// where the previous call's eight stopped.
|
|
//
|
|
// melspectrogram.onnx returns N/160-3 frames for N samples, and frame i
|
|
// covers [i*160, i*160+400). With 480 samples of history the buffer is
|
|
// 1760 samples, which is 8 frames, and the oldest of them starts one hop
|
|
// after the newest of the previous call. Less history leaves a gap: the
|
|
// first frames of a bare chunk would be computed against silence.
|
|
melContext = 480
|
|
|
|
// chunkSamples — audio per embedding step, 80ms.
|
|
chunkSamples = embedStride * melHop
|
|
)
|
|
|
|
// wakeModels holds the three ONNX sessions. It runs on CPU threads beside
|
|
// silero and never touches the GPU. That is a rule, not a result: a wake word
|
|
// that waits on card admission is not a wake word.
|
|
type wakeModels struct {
|
|
mel *ort.DynamicAdvancedSession
|
|
emb *ort.DynamicAdvancedSession
|
|
head *ort.DynamicAdvancedSession
|
|
}
|
|
|
|
// newWakeModels loads all three. melPath and embedPath are openWakeWord's
|
|
// frozen feature models; headPath is the keyword head trained for "Мэйвен".
|
|
func newWakeModels(melPath, embedPath, headPath, libPath string) (*wakeModels, error) {
|
|
if !ort.IsInitialized() {
|
|
if libPath != "" {
|
|
ort.SetSharedLibraryPath(libPath)
|
|
}
|
|
if err := ort.InitializeEnvironment(); err != nil {
|
|
return nil, fmt.Errorf("wake word: onnx runtime: %w", err)
|
|
}
|
|
}
|
|
// One thread per session, not the default of every core. Measured on
|
|
// workpc: the default took mavwaked from 68% of one core to 335% of
|
|
// three, for three graphs that each run in well under 80ms single
|
|
// threaded. An always-on gate that eats a quarter of the workstation is
|
|
// not a gate he will leave running.
|
|
opts, err := ort.NewSessionOptions()
|
|
if err != nil {
|
|
return nil, fmt.Errorf("wake word: session options: %w", err)
|
|
}
|
|
defer opts.Destroy()
|
|
if err := opts.SetIntraOpNumThreads(1); err != nil {
|
|
return nil, fmt.Errorf("wake word: intra-op threads: %w", err)
|
|
}
|
|
if err := opts.SetInterOpNumThreads(1); err != nil {
|
|
return nil, fmt.Errorf("wake word: inter-op threads: %w", err)
|
|
}
|
|
open := func(p string, in, out []string) (*ort.DynamicAdvancedSession, error) {
|
|
s, err := ort.NewDynamicAdvancedSession(p, in, out, opts)
|
|
if err != nil {
|
|
return nil, fmt.Errorf("wake word: load %s: %w", p, err)
|
|
}
|
|
return s, nil
|
|
}
|
|
m := &wakeModels{}
|
|
if m.mel, err = open(melPath, []string{"input"}, []string{"output"}); err != nil {
|
|
return nil, err
|
|
}
|
|
if m.emb, err = open(embedPath, []string{"input_1"}, []string{"conv2d_19"}); err != nil {
|
|
m.Close()
|
|
return nil, err
|
|
}
|
|
if m.head, err = open(headPath, []string{"embeddings"}, []string{"score"}); err != nil {
|
|
m.Close()
|
|
return nil, err
|
|
}
|
|
return m, nil
|
|
}
|
|
|
|
// Close releases the three sessions.
|
|
func (m *wakeModels) Close() {
|
|
if m == nil {
|
|
return
|
|
}
|
|
for _, s := range []*ort.DynamicAdvancedSession{m.mel, m.emb, m.head} {
|
|
if s != nil {
|
|
s.Destroy()
|
|
}
|
|
}
|
|
m.mel, m.emb, m.head = nil, nil, nil
|
|
}
|
|
|
|
// melFrames runs one buffer of samples and returns the mel frames it yielded.
|
|
func (m *wakeModels) melFrames(buf []float32) ([][melBins]float32, error) {
|
|
in, err := ort.NewTensor(ort.NewShape(1, int64(len(buf))), buf)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
defer in.Destroy()
|
|
|
|
n := int64(len(buf)/melHop - 3)
|
|
if n < 1 {
|
|
return nil, fmt.Errorf("wake word: %d samples yield no mel frames", len(buf))
|
|
}
|
|
out, err := ort.NewEmptyTensor[float32](ort.NewShape(1, 1, n, melBins))
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
defer out.Destroy()
|
|
|
|
if err := m.mel.Run([]ort.Value{in}, []ort.Value{out}); err != nil {
|
|
return nil, err
|
|
}
|
|
data := out.GetData()
|
|
frames := make([][melBins]float32, n)
|
|
for i := range frames {
|
|
for j := 0; j < melBins; j++ {
|
|
// The scaling openWakeWord applies between the two feature
|
|
// models, and the head was fitted on its output.
|
|
frames[i][j] = data[i*melBins+j]/10.0 + 2.0
|
|
}
|
|
}
|
|
return frames, nil
|
|
}
|
|
|
|
// embedding runs embedFrames mel frames through the frozen embedder.
|
|
func (m *wakeModels) embedding(mels [][melBins]float32) ([embedDim]float32, error) {
|
|
var e [embedDim]float32
|
|
flat := make([]float32, 0, embedFrames*melBins)
|
|
for _, f := range mels {
|
|
flat = append(flat, f[:]...)
|
|
}
|
|
in, err := ort.NewTensor(ort.NewShape(1, embedFrames, melBins, 1), flat)
|
|
if err != nil {
|
|
return e, err
|
|
}
|
|
defer in.Destroy()
|
|
out, err := ort.NewEmptyTensor[float32](ort.NewShape(1, 1, 1, embedDim))
|
|
if err != nil {
|
|
return e, err
|
|
}
|
|
defer out.Destroy()
|
|
if err := m.emb.Run([]ort.Value{in}, []ort.Value{out}); err != nil {
|
|
return e, err
|
|
}
|
|
copy(e[:], out.GetData())
|
|
return e, nil
|
|
}
|
|
|
|
// score runs the trained head over headWindow embeddings.
|
|
func (m *wakeModels) score(embeds [][embedDim]float32) (float64, error) {
|
|
flat := make([]float32, 0, headWindow*embedDim)
|
|
for _, e := range embeds {
|
|
flat = append(flat, e[:]...)
|
|
}
|
|
in, err := ort.NewTensor(ort.NewShape(1, headWindow, embedDim), flat)
|
|
if err != nil {
|
|
return 0, err
|
|
}
|
|
defer in.Destroy()
|
|
out, err := ort.NewEmptyTensor[float32](ort.NewShape(1, 1))
|
|
if err != nil {
|
|
return 0, err
|
|
}
|
|
defer out.Destroy()
|
|
if err := m.head.Run([]ort.Value{in}, []ort.Value{out}); err != nil {
|
|
return 0, err
|
|
}
|
|
return float64(out.GetData()[0]), nil
|
|
}
|