Add TorchScript/JIT module support for model inference

go.mod

@@ -1,6 +1,6 @@
 module github.com/lwch/gotorch
 
-go 1.20
+go 1.25
 
 require (
 	github.com/lwch/logging v1.1.3

internal/torch/api.h

@@ -3,12 +3,14 @@
 
 #ifdef __cplusplus
 #include <torch/torch.h>
+#include <torch/script.h>
 extern "C"
 {
 
     typedef torch::Tensor *tensor;
     typedef torch::optim::Optimizer *optimizer;
     typedef torch::nn::Module *module;
+    typedef torch::jit::script::Module *jit_module;
 
     struct _optimizer_state
     {
@@ -19,6 +21,7 @@ extern "C"
 typedef void *tensor;
 typedef void *optimizer;
 typedef void *module;
+typedef void *jit_module;
 typedef void *optimizer_state;
 #endif
 

internal/torch/exception.hpp

@@ -129,6 +129,24 @@ optimizer_state auto_catch_optimizer_state(Function f, char **err)
     return 0;
 }
 
+template <typename Function>
+jit_module auto_catch_jit_module(Function f, char **err)
+{
+    try
+    {
+        return f();
+    }
+    catch (const torch::Error &e)
+    {
+        *err = strdup(e.msg().c_str());
+    }
+    catch (const std::exception &e)
+    {
+        *err = strdup(e.what());
+    }
+    return nullptr;
+}
+
 #endif
 
 #endif // __GOTORCH_EXCEPTION_H__

internal/torch/jit.go

@@ -0,0 +1,107 @@
+package torch
+
+// #cgo noescape jit_forward
+// #cgo noescape jit_forward_multi
+// #cgo nocallback jit_forward
+// #cgo nocallback jit_forward_multi
+// #cgo nocallback jit_eval
+// #cgo nocallback jit_train
+// #cgo nocallback jit_free
+// #include <stdlib.h>
+// #include "jit.h"
+import "C"
+import (
+	"unsafe"
+
+	"github.com/lwch/gotorch/consts"
+)
+
+// JitModule is the low-level handle to a TorchScript module
+type JitModule C.jit_module
+
+// JitLoad loads a TorchScript model from file (CPU)
+func JitLoad(path string) JitModule {
+	cpath := C.CString(path)
+	defer C.free(unsafe.Pointer(cpath))
+
+	var err *C.char
+	m := C.jit_load(&err, cpath)
+	if err != nil {
+		defer C.free(unsafe.Pointer(err))
+		panic(C.GoString(err))
+	}
+	return JitModule(m)
+}
+
+// JitLoadToDevice loads a TorchScript model to specified device
+func JitLoadToDevice(path string, device consts.DeviceType) JitModule {
+	cpath := C.CString(path)
+	defer C.free(unsafe.Pointer(cpath))
+
+	var err *C.char
+	m := C.jit_load_to_device(&err, cpath, C.int8_t(device))
+	if err != nil {
+		defer C.free(unsafe.Pointer(err))
+		panic(C.GoString(err))
+	}
+	return JitModule(m)
+}
+
+// JitForward runs forward pass with single output
+func JitForward(m JitModule, input Tensor) Tensor {
+	var err *C.char
+	out := C.jit_forward(&err, C.jit_module(m), C.tensor(input))
+	if err != nil {
+		defer C.free(unsafe.Pointer(err))
+		panic(C.GoString(err))
+	}
+	return Tensor(out)
+}
+
+// JitForwardMulti runs forward pass returning multiple outputs
+func JitForwardMulti(m JitModule, input Tensor, numOutputs int) []Tensor {
+	if numOutputs <= 0 {
+		return nil
+	}
+
+	outputs := make([]C.tensor, numOutputs)
+
+	var err *C.char
+	actual := C.jit_forward_multi(&err, C.jit_module(m), C.tensor(input),
+		&outputs[0], C.size_t(numOutputs))
+	if err != nil {
+		defer C.free(unsafe.Pointer(err))
+		panic(C.GoString(err))
+	}
+
+	result := make([]Tensor, actual)
+	for i := range int(actual) {
+		result[i] = Tensor(outputs[i])
+	}
+	return result
+}
+
+// JitToDevice moves module to specified device
+func JitToDevice(m JitModule, device consts.DeviceType) {
+	var err *C.char
+	C.jit_to_device(&err, C.jit_module(m), C.int8_t(device))
+	if err != nil {
+		defer C.free(unsafe.Pointer(err))
+		panic(C.GoString(err))
+	}
+}
+
+// JitEval sets module to evaluation mode
+func JitEval(m JitModule) {
+	C.jit_eval(C.jit_module(m))
+}
+
+// JitTrain sets module to training mode
+func JitTrain(m JitModule) {
+	C.jit_train(C.jit_module(m))
+}
+
+// JitFree releases module resources
+func JitFree(m JitModule) {
+	C.jit_free(C.jit_module(m))
+}

internal/torch/jit.h

@@ -0,0 +1,44 @@
+#ifndef __GOTORCH_JIT_H__
+#define __GOTORCH_JIT_H__
+
+#include <stddef.h>
+#include <stdint.h>
+#include "api.h"
+
+#ifdef __cplusplus
+extern "C"
+{
+#endif
+
+    // Load TorchScript model from file (defaults to CPU)
+    GOTORCH_API jit_module jit_load(char **err, const char *path);
+
+    // Load TorchScript model to specified device (0=CPU, 1=CUDA)
+    GOTORCH_API jit_module jit_load_to_device(char **err, const char *path, int8_t device);
+
+    // Forward pass with single output (returns first tensor if tuple)
+    GOTORCH_API tensor jit_forward(char **err, jit_module m, tensor input);
+
+    // Forward pass with multiple outputs (for models returning tuples)
+    // out_tensors: pre-allocated array, out_count: array size
+    // Returns: actual number of outputs written
+    GOTORCH_API size_t jit_forward_multi(char **err, jit_module m, tensor input,
+                                          tensor *out_tensors, size_t out_count);
+
+    // Move module to device
+    GOTORCH_API void jit_to_device(char **err, jit_module m, int8_t device);
+
+    // Set evaluation mode (disables dropout, batch norm updates)
+    GOTORCH_API void jit_eval(jit_module m);
+
+    // Set training mode
+    GOTORCH_API void jit_train(jit_module m);
+
+    // Free module resources
+    GOTORCH_API void jit_free(jit_module m);
+
+#ifdef __cplusplus
+}
+#endif
+
+#endif // __GOTORCH_JIT_H__

jit/jit.go

@@ -0,0 +1,172 @@
+// Package jit provides TorchScript model loading and inference capabilities.
+package jit
+
+import (
+	"errors"
+	"fmt"
+	"iter"
+	"runtime"
+	"sync"
+	"sync/atomic"
+
+	"github.com/lwch/gotorch/consts"
+	"github.com/lwch/gotorch/internal/torch"
+	"github.com/lwch/gotorch/tensor"
+)
+
+// ErrModuleClosed is returned when operations are attempted on a closed module.
+var ErrModuleClosed = errors.New("module has been closed")
+
+// moduleHandle wraps the JIT module with once-only cleanup semantics.
+type moduleHandle struct {
+	m      torch.JitModule
+	once   sync.Once
+	closed atomic.Bool
+}
+
+func (h *moduleHandle) free() {
+	h.once.Do(func() {
+		h.closed.Store(true)
+		torch.JitFree(h.m)
+	})
+}
+
+func (h *moduleHandle) isClosed() bool {
+	return h.closed.Load()
+}
+
+// Module represents a loaded TorchScript model.
+type Module struct {
+	handle *moduleHandle
+}
+
+// Load loads a TorchScript model from file to CPU.
+func Load(path string) (m *Module, err error) {
+	return LoadToDevice(path, consts.KCPU)
+}
+
+// LoadToDevice loads a TorchScript model to the specified device.
+func LoadToDevice(path string, device consts.DeviceType) (m *Module, err error) {
+	defer func() {
+		if r := recover(); r != nil {
+			m = nil
+			err = fmt.Errorf("failed to load model %q: %v", path, r)
+		}
+	}()
+
+	jm := torch.JitLoadToDevice(path, device)
+	handle := &moduleHandle{m: jm}
+	module := &Module{handle: handle}
+
+	// Go 1.24+: Use AddCleanup instead of SetFinalizer
+	// Benefits: no cycle leaks, multiple cleanups allowed, works with interior pointers
+	// The sync.Once in moduleHandle ensures free is called exactly once
+	runtime.AddCleanup(module, func(h *moduleHandle) {
+		h.free()
+	}, handle)
+
+	return module, nil
+}
+
+// Close explicitly releases module resources.
+// The module should not be used after calling Close.
+// Safe to call multiple times.
+func (m *Module) Close() {
+	if m != nil && m.handle != nil {
+		m.handle.free()
+	}
+}
+
+// Forward runs inference and returns a single output tensor.
+// For models returning multiple outputs, this returns only the first.
+func (m *Module) Forward(input *tensor.Tensor) (out *tensor.Tensor, err error) {
+	if m == nil || m.handle == nil {
+		return nil, ErrModuleClosed
+	}
+	if m.handle.isClosed() {
+		return nil, ErrModuleClosed
+	}
+
+	defer func() {
+		if r := recover(); r != nil {
+			out = nil
+			err = fmt.Errorf("forward pass failed: %v", r)
+		}
+	}()
+
+	result := torch.JitForward(m.handle.m, input.Tensor())
+	return tensor.New(result), nil
+}
+
+// ForwardMulti runs inference and returns multiple output tensors.
+// This is useful for models like BirdNET v3.0 that return (embeddings, predictions).
+func (m *Module) ForwardMulti(input *tensor.Tensor, numOutputs int) (out []*tensor.Tensor, err error) {
+	if m == nil || m.handle == nil {
+		return nil, ErrModuleClosed
+	}
+	if m.handle.isClosed() {
+		return nil, ErrModuleClosed
+	}
+
+	defer func() {
+		if r := recover(); r != nil {
+			out = nil
+			err = fmt.Errorf("forward pass failed: %v", r)
+		}
+	}()
+
+	outputs := torch.JitForwardMulti(m.handle.m, input.Tensor(), numOutputs)
+	result := make([]*tensor.Tensor, len(outputs))
+	for i, t := range outputs {
+		result[i] = tensor.New(t)
+	}
+	return result, nil
+}
+
+// ToDevice moves the module to the specified device.
+func (m *Module) ToDevice(device consts.DeviceType) (err error) {
+	if m == nil || m.handle == nil {
+		return ErrModuleClosed
+	}
+	if m.handle.isClosed() {
+		return ErrModuleClosed
+	}
+
+	defer func() {
+		if r := recover(); r != nil {
+			err = fmt.Errorf("failed to move module to device: %v", r)
+		}
+	}()
+
+	torch.JitToDevice(m.handle.m, device)
+	return nil
+}
+
+// Eval sets the module to evaluation mode.
+// This disables dropout and batch normalization updates.
+func (m *Module) Eval() {
+	if m == nil || m.handle == nil || m.handle.isClosed() {
+		return
+	}
+	torch.JitEval(m.handle.m)
+}
+
+// Train sets the module to training mode.
+func (m *Module) Train() {
+	if m == nil || m.handle == nil || m.handle.isClosed() {
+		return
+	}
+	torch.JitTrain(m.handle.m)
+}
+
+// Outputs returns an iterator over the output tensors from ForwardMulti.
+// This leverages Go 1.23+ iterator support for cleaner code.
+func Outputs(tensors []*tensor.Tensor) iter.Seq2[int, *tensor.Tensor] {
+	return func(yield func(int, *tensor.Tensor) bool) {
+		for i, t := range tensors {
+			if !yield(i, t) {
+				return
+			}
+		}
+	}
+}