chore(v2): vendor dependencies for offline/China builds
go mod vendor pins onnxruntime_go v1.12.1, Gio and the rest into v2/vendor so go run/build work without hitting proxy.golang.org (blocked/slow in China). Verified: CGO_ENABLED=1 go build -mod=vendor ./internal/spike and GOOS=windows go build -mod=vendor ./internal/ui both pass. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
This commit is contained in:
Claude Opus 4.8
@@ -0,0 +1,3 @@
|
||||
onnxruntime_c_api.h linguist-vendored
|
||||
onnxruntime_training_c_api.h linguist-vendored
|
||||
|
||||
+138
@@ -0,0 +1,138 @@
|
||||
Contribution Guidelines
|
||||
=======================
|
||||
|
||||
This library began as a personal project, and is primarily still maintained as
|
||||
such. The following list of guidelines is not necessarily exhaustive, and,
|
||||
ultimately, any contribution is subject to the maintainer's discretion. That
|
||||
being said, contributions are welcome, and most recent new features have been
|
||||
added by users who need them!
|
||||
|
||||
Coding Style
|
||||
------------
|
||||
|
||||
- Go code must be formatted using the official `gofmt` tool.
|
||||
|
||||
- C code should adhere to the portions of Google's C++ style guide that
|
||||
apply to C.
|
||||
|
||||
- If at all possible, any Go or C code should have at most 80 character lines.
|
||||
(This may not be enforced very strictly.)
|
||||
|
||||
- Purely stylistic changes are unlikely to be accepted. Instead, the
|
||||
maintainer or other contributers may make small stylistic adjustments to
|
||||
surrounding code as part of other contributions.
|
||||
|
||||
- Attempt to mimic the existing style of the surrounding code.
|
||||
|
||||
|
||||
Documentation
|
||||
-------------
|
||||
|
||||
- All Go types, public-facing functions, and nontrivial internal functions
|
||||
must include a comment on their intended usage, to be parsed by godoc.
|
||||
|
||||
- As per the google C++ style guide, all C functions must be documented with a
|
||||
comment as well. If a C function is defined in a header file, the comment
|
||||
should appear with the definition in the header. If it's a static function
|
||||
in a `.c` file, the comment should appear with the function definition.
|
||||
|
||||
|
||||
Tests
|
||||
-----
|
||||
|
||||
- All new features and bugfixes must include a basic unit test (in
|
||||
`onnxruntime_test.go` or `onnxruntime_training_test.go`) to serve as a
|
||||
sanity check.
|
||||
|
||||
- If a test is for a platform-dependent or execution-provider-dependent
|
||||
feature, the test must be skipped if run on an unsupported system.
|
||||
|
||||
- No tests should panic. Always check errors and fail rather than allowing
|
||||
tests to panic.
|
||||
|
||||
- Every change must ensure that `go test -v -bench=.` passes on every
|
||||
supported platform.
|
||||
|
||||
- Every test failure should be accompanied by a message containing the reason,
|
||||
either using `t.Logf()`, `t.Errorf()`, or `t.Fatalf()`.
|
||||
|
||||
|
||||
Adding New Files
|
||||
----------------
|
||||
|
||||
- Apart from testing data, try not to add new source files.
|
||||
|
||||
- Do not add third-party code or headers. The only exceptions for now are
|
||||
`onnxruntime_c_api.h` and `onnxruntime_training_c_api.h`.
|
||||
|
||||
- No C++ at all. Developing Go-to-C wrappers is annoying enough as it is.
|
||||
|
||||
- Do not add any new `onnxruntime` shared libraries under `test_data`. I know
|
||||
there are additional platforms that would be nice to include (such as
|
||||
`x86_64` Linux), but I do not want this project turning into an unofficial
|
||||
distribution channel for onnxruntime libraries. It also clogs up the git
|
||||
repo with large files, and increases the size of the history every time
|
||||
these files are updated. The libraries that are included were only intended
|
||||
to allow a majority of users to run `go test -v -bench=.` without further
|
||||
setup or modification. Currently: amd64 Windows, arm64 Linux (I wish I
|
||||
hadn't included this!), arm64 osx, and amd64 osx. All other users must set
|
||||
the `ONNXRUNTIME_SHARED_LIBRARY_PATH` environment variable to a valid path
|
||||
to the correct `onnxruntime` shared library file prior to running tests.
|
||||
|
||||
- If you need to add a .onnx file for a test, place both the .onnx file
|
||||
_and_ the script used to generate it into `test_data/`.
|
||||
|
||||
- Keep any testing .onnx files as small as possible.
|
||||
|
||||
- Without a good reason (i.e., implementing an entire class of APIs such as
|
||||
training), avoid adding new Go files---just add to `onnxruntime_go.go`.
|
||||
|
||||
|
||||
Dependencies
|
||||
------------
|
||||
|
||||
- Avoid Go or C dependencies outside of the language's standard libraries.
|
||||
This package currently does not depend on any third-party Go modules, and
|
||||
it would be great to keep it this way.
|
||||
|
||||
- Python scripts within `test_data/` can use whatever dependencies they need,
|
||||
because end users should not be required to run the python files, and the
|
||||
`.onnx` file they produce should already be included.
|
||||
|
||||
|
||||
C-Specific Stuff
|
||||
----------------
|
||||
|
||||
- Minimize Go management of C-allocated memory as much as possible. For
|
||||
example, see the `convertORTString` function on `onnxruntime_go.go`, which
|
||||
copies a C-allocated string into a garbage-collected go `string`.
|
||||
|
||||
- If you need to use a `OrtAllocator` in onnxruntime's C API, always use the
|
||||
default `OrtAllocator` returned by
|
||||
`ort_api->GetAllocatorWithDefaultOptions()`.
|
||||
|
||||
- ONNXRuntime APIs requiring file paths typically use `ORTCHAR_T*`
|
||||
strings. On Linux/OSX/etc, these should be UTF-8, but on Windows they will
|
||||
be wide-character strings. (Our tricks with `#include` to make them look
|
||||
like `char*` to C code even on Windows, but the DLL still expects a
|
||||
`wchar_t*`.) The important takeaway: when passing `ORTCHAR_T*`
|
||||
values to the onnxruntime C API, use the `createOrtCharString(...)`
|
||||
function. It converts a Go string to a C string, but unlike `C.CString`, it
|
||||
will do UTF8 to UTF16 conversion on Windows. (On Linux, it simply wraps
|
||||
`C.CString`.)
|
||||
|
||||
|
||||
A Few Notes on Organization
|
||||
---------------------------
|
||||
|
||||
- The `onnxruntime` C API uses a struct containing function pointers. Cgo
|
||||
can't directly invoke functions via pointers, so `onnxruntime_wrapper.c`
|
||||
(along with the associated header file) are used to provide top-level C
|
||||
functions that call the function pointers within the `OrtApi` struct.
|
||||
|
||||
- Linux and OSX use `dlopen` to load the onnxruntime shared library, but this
|
||||
isn't possible on Windows, which instead can use the `syscall.LoadLibrary()`
|
||||
function from Go's standard library. This different behavior is locked
|
||||
behind build constraints in `setup_env.go` and `setup_env_windows.go`,
|
||||
respectively.
|
||||
|
||||
+19
@@ -0,0 +1,19 @@
|
||||
Copyright (c) 2023 Nathan Otterness
|
||||
|
||||
Permission is hereby granted, free of charge, to any person obtaining a copy
|
||||
of this software and associated documentation files (the "Software"), to deal
|
||||
in the Software without restriction, including without limitation the rights
|
||||
to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
|
||||
copies of the Software, and to permit persons to whom the Software is
|
||||
furnished to do so, subject to the following conditions:
|
||||
|
||||
The above copyright notice and this permission notice shall be included in all
|
||||
copies or substantial portions of the Software.
|
||||
|
||||
THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
|
||||
IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
|
||||
FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
|
||||
AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
|
||||
LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
|
||||
OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
|
||||
SOFTWARE.
|
||||
+216
@@ -0,0 +1,216 @@
|
||||
Cross-Platform `onnxruntime` Wrapper for Go
|
||||
===========================================
|
||||
|
||||
About
|
||||
-----
|
||||
|
||||
This library seeks to provide an interface for loading and executing neural
|
||||
networks from Go(lang) code, while remaining as simple to use as possible.
|
||||
|
||||
A few example applications using this library can be found in the
|
||||
[`onnxruntime_go_examples` repository](https://github.com/yalue/onnxruntime_go_examples).
|
||||
|
||||
The [onnxruntime](https://github.com/microsoft/onnxruntime) library provides a
|
||||
way to load and execute ONNX-format neural networks, though the library
|
||||
primarily supports C and C++ APIs. Several efforts exist to have written
|
||||
Go(lang) wrappers for the `onnxruntime` library, but as far as I can tell, none
|
||||
of these existing Go wrappers support Windows. This is due to the fact that
|
||||
Microsoft's `onnxruntime` library assumes the user will be using the MSVC
|
||||
compiler on Windows systems, while CGo on Windows requires using Mingw.
|
||||
|
||||
This wrapper works around the issues by manually loading the `onnxruntime`
|
||||
shared library, removing any dependency on the `onnxruntime` source code beyond
|
||||
the header files. Naturally, this approach works equally well on non-Windows
|
||||
systems.
|
||||
|
||||
Additionally, this library uses Go's recent addition of generics to support
|
||||
multiple Tensor data types; see the `NewTensor` or `NewEmptyTensor` functions.
|
||||
|
||||
**IMPORTANT:** As of onnxruntime_go v1.12.0 or above, for CUDA acceleration we
|
||||
now require the use of CUDA 12.x and CuDNN 9.x (as required by onnxruntime
|
||||
v1.19.0+). Those wishing to stay on CUDA 11.8 should remain on onnxruntime_go
|
||||
v1.11.0 or below.
|
||||
|
||||
Note on onnxruntime Library Versions
|
||||
------------------------------------
|
||||
|
||||
At the time of writing, this library uses version 1.19.0 of the onnxruntime
|
||||
C API headers. So, it will probably only work with version 1.19.0 of the
|
||||
onnxruntime shared libraries, as well. If you need to use a different version,
|
||||
or if I get behind on updating this repository, updating or changing the
|
||||
onnxruntime version should be fairly easy:
|
||||
|
||||
1. Replace the `onnxruntime_c_api.h` file with the version corresponding to
|
||||
the onnxruntime version you wish to use.
|
||||
|
||||
2. Replace the `test_data/onnxruntime.dll` (or `test_data/onnxruntime*.so`)
|
||||
file with the version corresponding to the onnxruntime version you wish to
|
||||
use.
|
||||
|
||||
3. (If you care about DirectML support) Verify that the entries in the
|
||||
`DummyOrtDMLAPI` struct in `onnxruntime_wrapper.c` match the order in which
|
||||
they appear in the `OrtDmlApi` struct from the `dml_provider_factory.h`
|
||||
header in the official repo. See the comment on this struct in
|
||||
`onnxruntime_wrapper.c` for more information.
|
||||
|
||||
Note that both the C API header and the shared library files are available to
|
||||
download from the releases page in the
|
||||
[official repo](https://github.com/microsoft/onnxruntime). Download the archive
|
||||
for the release you want to use, and extract it. The header file is located in
|
||||
the "include" subdirectory, and the shared library will be located in the "lib"
|
||||
subdirectory. (On Linux systems, you'll need the version of the .so with the
|
||||
appended version numbers, e.g., `libonnxruntime.so.1.19.0`, and _not_ the
|
||||
`libonnxruntime.so`, which is just a symbolic link.) The archive will contain
|
||||
several other files containing C++ headers, debug symbols, and so on, but you
|
||||
shouldn't need anything other than the single onnxruntime shared library and
|
||||
`onnxruntime_c_api.h`. (The exception is if you're wanting to enable GPU
|
||||
support, where you may need other shared-library files, such as
|
||||
`execution_providers_cuda.dll` and `execution_providers_shared.dll` on Windows.)
|
||||
|
||||
|
||||
Requirements
|
||||
------------
|
||||
|
||||
To use this library, you'll need a version of Go with cgo support. If you are
|
||||
not using an amd64 version of Windows or Linux (or if you want to provide your
|
||||
own library for some other reason), you simply need to provide the correct path
|
||||
to the shared library when initializing the wrapper. This is seen in the first
|
||||
few lines of the following example.
|
||||
|
||||
Note that if you want to use CUDA, you'll need to be using a version of the
|
||||
onnxruntime shared library with CUDA support, as well as be using a CUDA
|
||||
version supported by the underlying version of your onnxruntime library. For
|
||||
example, version 1.19.0 of the onnxruntime library only supports CUDA versions
|
||||
12.x. See
|
||||
[the onnxruntime CUDA support documentation](https://onnxruntime.ai/docs/execution-providers/CUDA-ExecutionProvider.html)
|
||||
for more specifics.
|
||||
|
||||
|
||||
Example Usage
|
||||
-------------
|
||||
|
||||
The full documentation can be found at [pkg.go.dev](https://pkg.go.dev/github.com/yalue/onnxruntime_go).
|
||||
|
||||
Additionally, several example command-line applications complete with necessary
|
||||
networks and data can be found in the
|
||||
[`onnxruntime_go_examples` repository](https://github.com/yalue/onnxruntime_go_examples).
|
||||
|
||||
The following example illustrates how this library can be used to load and run
|
||||
an ONNX network taking a single input tensor and producing a single output
|
||||
tensor, both of which contain 32-bit floating point values. Note that error
|
||||
handling is omitted; each of the functions returns an err value, which will be
|
||||
non-nil in the case of failure.
|
||||
|
||||
```go
|
||||
import (
|
||||
"fmt"
|
||||
ort "github.com/yalue/onnxruntime_go"
|
||||
"os"
|
||||
)
|
||||
|
||||
func main() {
|
||||
// This line _may_ be optional; by default the library will try to load
|
||||
// "onnxruntime.dll" on Windows, and "onnxruntime.so" on any other system.
|
||||
// For stability, it is probably a good idea to always set this explicitly.
|
||||
ort.SetSharedLibraryPath("path/to/onnxruntime.so")
|
||||
|
||||
err := ort.InitializeEnvironment()
|
||||
if err != nil {
|
||||
panic(err)
|
||||
}
|
||||
defer ort.DestroyEnvironment()
|
||||
|
||||
// For a slight performance boost and convenience when re-using existing
|
||||
// tensors, this library expects the user to create all input and output
|
||||
// tensors prior to creating the session. If this isn't ideal for your use
|
||||
// case, see the DynamicAdvancedSession type in the documnentation, which
|
||||
// allows input and output tensors to be specified when calling Run()
|
||||
// rather than when initializing a session.
|
||||
inputData := []float32{0.0, 0.1, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7, 0.8, 0.9}
|
||||
inputShape := ort.NewShape(2, 5)
|
||||
inputTensor, err := ort.NewTensor(inputShape, inputData)
|
||||
defer inputTensor.Destroy()
|
||||
// This hypothetical network maps a 2x5 input -> 2x3x4 output.
|
||||
outputShape := ort.NewShape(2, 3, 4)
|
||||
outputTensor, err := ort.NewEmptyTensor[float32](outputShape)
|
||||
defer outputTensor.Destroy()
|
||||
|
||||
session, err := ort.NewAdvancedSession("path/to/network.onnx",
|
||||
[]string{"Input 1 Name"}, []string{"Output 1 Name"},
|
||||
[]ort.Value{inputTensor}, []ort.Value{outputTensor}, nil)
|
||||
defer session.Destroy()
|
||||
|
||||
// Calling Run() will run the network, reading the current contents of the
|
||||
// input tensors and modifying the contents of the output tensors.
|
||||
err = session.Run()
|
||||
|
||||
// Get a slice view of the output tensor's data.
|
||||
outputData := outputTensor.GetData()
|
||||
|
||||
// If you want to run the network on a different input, all you need to do
|
||||
// is modify the input tensor data (available via inputTensor.GetData())
|
||||
// and call Run() again.
|
||||
|
||||
// ...
|
||||
}
|
||||
```
|
||||
|
||||
|
||||
Deprecated APIs
|
||||
---------------
|
||||
|
||||
Older versions of this library used a typed `Session[T]` struct to keep track
|
||||
of sessions. In retrospect, associating type parameters with Sessions was
|
||||
unnecessary, and the `AdvancedSession` type, along with its associated APIs,
|
||||
was added to rectify this mistake. For backwards compatibility, the old typed
|
||||
`Session[T]` and `DynamicSession[T]` types are still included and unlikely to
|
||||
be removed. However, they now delegate their functionality to
|
||||
`AdvancedSession` internally. New code should always favor using
|
||||
`AdvancedSession` directly.
|
||||
|
||||
|
||||
Running Tests and System Compatibility for Testing
|
||||
--------------------------------------------------
|
||||
|
||||
Navigate to this directory and run `go test -v`, or optionally
|
||||
`go test -v -bench=.`. All tests should pass; tests relating to CUDA or other
|
||||
accelerator support will be skipped on systems or onnxruntime builds that don't
|
||||
support them.
|
||||
|
||||
Currently, this repository includes a copy of the onnxruntime shared libraries
|
||||
for a few systems, including AMD64 windows, ARM64 Linux, and ARM64 darwin.
|
||||
These should allow tests to pass on those systems without users needing to copy
|
||||
additional libraries beyond cloning this repository. In the future, however,
|
||||
this may change if support for more systems are added or removed.
|
||||
|
||||
You may want to use a different version of the `onnxruntime` shared library for
|
||||
a couple reasons. In particular:
|
||||
|
||||
1. The included shared library copies do not include support for CUDA or other
|
||||
accelerated execution providers, so CUDA-related tests will always be
|
||||
skipped if you use the default libraries in this repo.
|
||||
|
||||
2. Many systems, including AMD64 and i386 Linux, and x86 osx, do not currently
|
||||
have shared libraries included in `test_data/` in the first place. (I would
|
||||
like to keep this directory, and the overall repo, smaller by keeping the
|
||||
number of shared libraries small.)
|
||||
|
||||
If these or other reasons apply to you, the test code will check the
|
||||
`ONNXRUNTIME_SHARED_LIBRARY_PATH` environment variable before attempting to
|
||||
load a library from `test_data/`. So, if you are using one of these systems or
|
||||
want accelerator-related tests to run, you should set the environment variable
|
||||
to the path to the onnxruntime shared library. Afterwards, `go test -v` should
|
||||
run and pass.
|
||||
|
||||
|
||||
Training API Support
|
||||
--------------------
|
||||
|
||||
This wrapper supports the onnxruntime training API on limited platforms. See
|
||||
the `NewTrainingSession` and associated data types or functions to use it. So
|
||||
far, the training API has only been tested on Linux, on `x86_64` architectures.
|
||||
|
||||
If you are not sure whether your platform or build of onnxruntime supports
|
||||
training, you can call `onnxruntime_go.IsTrainingSupported()`, which will
|
||||
return `true` if training is supported on your system.
|
||||
|
||||
+173
@@ -0,0 +1,173 @@
|
||||
package onnxruntime_go
|
||||
|
||||
// This file contains Session types that we maintain for compatibility
|
||||
// purposes; the main onnxruntime_go.go file is dedicated to AdvancedSession
|
||||
// now.
|
||||
|
||||
import (
|
||||
"fmt"
|
||||
"os"
|
||||
)
|
||||
|
||||
// #include "onnxruntime_wrapper.h"
|
||||
import "C"
|
||||
|
||||
// This type of session is for ONNX networks with the same input and output
|
||||
// data types.
|
||||
//
|
||||
// NOTE: This type was written with a type parameter despite the fact that a
|
||||
// type parameter is not necessary for any of its underlying implementation,
|
||||
// which is a mistake in retrospect. It is preserved only for compatibility
|
||||
// with older code, and new users should almost certainly be using an
|
||||
// AdvancedSession instead.
|
||||
//
|
||||
// Using an AdvancedSession struct should be easier, and supports arbitrary
|
||||
// combination of input and output tensor data types as well as more options.
|
||||
type Session[T TensorData] struct {
|
||||
// We now delegate all of the implementation to an AdvancedSession here.
|
||||
s *AdvancedSession
|
||||
}
|
||||
|
||||
// Similar to Session, but does not require the specification of the input
|
||||
// and output shapes at session creation time, and allows for input and output
|
||||
// tensors to have different types. This allows for fully dynamic input to the
|
||||
// onnx model.
|
||||
//
|
||||
// NOTE: As with Session[T], new users should probably be using
|
||||
// DynamicAdvancedSession in the future.
|
||||
type DynamicSession[In TensorData, Out TensorData] struct {
|
||||
s *DynamicAdvancedSession
|
||||
}
|
||||
|
||||
// The same as NewSession, but takes a slice of bytes containing the .onnx
|
||||
// network rather than a file path.
|
||||
func NewSessionWithONNXData[T TensorData](onnxData []byte, inputNames,
|
||||
outputNames []string, inputs, outputs []*Tensor[T]) (*Session[T], error) {
|
||||
// Unfortunately, a slice of pointers that satisfy an interface don't count
|
||||
// as a slice of interfaces (at least, as I write this), so we'll make the
|
||||
// conversion here.
|
||||
tmpInputs := make([]Value, len(inputs))
|
||||
tmpOutputs := make([]Value, len(outputs))
|
||||
for i, t := range inputs {
|
||||
tmpInputs[i] = t
|
||||
}
|
||||
for i, t := range outputs {
|
||||
tmpOutputs[i] = t
|
||||
}
|
||||
s, e := NewAdvancedSessionWithONNXData(onnxData, inputNames, outputNames,
|
||||
tmpInputs, tmpOutputs, nil)
|
||||
if e != nil {
|
||||
return nil, e
|
||||
}
|
||||
return &Session[T]{
|
||||
s: s,
|
||||
}, nil
|
||||
}
|
||||
|
||||
// Similar to NewSessionWithOnnxData, but for dynamic sessions.
|
||||
func NewDynamicSessionWithONNXData[in TensorData, out TensorData](onnxData []byte,
|
||||
inputNames, outputNames []string) (*DynamicSession[in, out], error) {
|
||||
s, e := NewDynamicAdvancedSessionWithONNXData(onnxData, inputNames,
|
||||
outputNames, nil)
|
||||
if e != nil {
|
||||
return nil, e
|
||||
}
|
||||
return &DynamicSession[in, out]{
|
||||
s: s,
|
||||
}, nil
|
||||
}
|
||||
|
||||
// Loads the ONNX network at the given path, and initializes a Session
|
||||
// instance. If this returns successfully, the caller must call Destroy() on
|
||||
// the returned session when it is no longer needed. We require the user to
|
||||
// provide the input and output tensors and names at this point, in order to
|
||||
// not need to re-allocate them every time Run() is called. The user instead
|
||||
// can just update or access the input/output tensor data after calling Run().
|
||||
// The input and output tensors MUST outlive this session, and calling
|
||||
// session.Destroy() will not destroy the input or output tensors.
|
||||
func NewSession[T TensorData](onnxFilePath string, inputNames,
|
||||
outputNames []string, inputs, outputs []*Tensor[T]) (*Session[T], error) {
|
||||
fileContent, e := os.ReadFile(onnxFilePath)
|
||||
if e != nil {
|
||||
return nil, fmt.Errorf("Error reading %s: %w", onnxFilePath, e)
|
||||
}
|
||||
|
||||
toReturn, e := NewSessionWithONNXData[T](fileContent, inputNames,
|
||||
outputNames, inputs, outputs)
|
||||
if e != nil {
|
||||
return nil, fmt.Errorf("Error creating session from %s: %w",
|
||||
onnxFilePath, e)
|
||||
}
|
||||
return toReturn, nil
|
||||
}
|
||||
|
||||
// Same as NewSession, but for dynamic sessions.
|
||||
func NewDynamicSession[in TensorData, out TensorData](onnxFilePath string,
|
||||
inputNames, outputNames []string) (*DynamicSession[in, out], error) {
|
||||
fileContent, e := os.ReadFile(onnxFilePath)
|
||||
if e != nil {
|
||||
return nil, fmt.Errorf("Error reading %s: %w", onnxFilePath, e)
|
||||
}
|
||||
|
||||
toReturn, e := NewDynamicSessionWithONNXData[in, out](fileContent,
|
||||
inputNames, outputNames)
|
||||
if e != nil {
|
||||
return nil, fmt.Errorf("Error creating session from %s: %w",
|
||||
onnxFilePath, e)
|
||||
}
|
||||
return toReturn, nil
|
||||
}
|
||||
|
||||
func (s *Session[_]) Destroy() error {
|
||||
return s.s.Destroy()
|
||||
}
|
||||
|
||||
func (s *DynamicSession[_, _]) Destroy() error {
|
||||
return s.s.Destroy()
|
||||
}
|
||||
|
||||
func (s *Session[T]) Run() error {
|
||||
return s.s.Run()
|
||||
}
|
||||
|
||||
// Unlike the non-dynamic equivalents, the DynamicSession's Run() function
|
||||
// takes a list of input and output tensors rather than requiring the tensors
|
||||
// to be specified at Session creation time. It is still the caller's
|
||||
// responsibility to create and Destroy all tensors passed to this function.
|
||||
func (s *DynamicSession[in, out]) Run(inputs []*Tensor[in],
|
||||
outputs []*Tensor[out]) error {
|
||||
if len(inputs) != len(s.s.s.inputNames) {
|
||||
return fmt.Errorf("The session specified %d input names, but Run() "+
|
||||
"was called with %d input tensors", len(s.s.s.inputNames),
|
||||
len(inputs))
|
||||
}
|
||||
if len(outputs) != len(s.s.s.outputNames) {
|
||||
return fmt.Errorf("The session specified %d output names, but Run() "+
|
||||
"was called with %d output tensors", len(s.s.s.outputNames),
|
||||
len(outputs))
|
||||
}
|
||||
inputValues := make([]*C.OrtValue, len(inputs))
|
||||
for i, v := range inputs {
|
||||
inputValues[i] = v.GetInternals().ortValue
|
||||
}
|
||||
outputValues := make([]*C.OrtValue, len(outputs))
|
||||
for i, v := range outputs {
|
||||
outputValues[i] = v.GetInternals().ortValue
|
||||
}
|
||||
|
||||
status := C.RunOrtSession(s.s.s.ortSession, &inputValues[0],
|
||||
&s.s.s.inputNames[0], C.int(len(inputs)), &outputValues[0],
|
||||
&s.s.s.outputNames[0], C.int(len(outputs)))
|
||||
if status != nil {
|
||||
return fmt.Errorf("Error running network: %w", statusToError(status))
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// This type alias is included to avoid breaking older code, where the inputs
|
||||
// and outputs to session.Run() were ArbitraryTensors rather than Values.
|
||||
type ArbitraryTensor = Value
|
||||
|
||||
// As with the ArbitraryTensor type, this type alias only exists to facilitate
|
||||
// renaming an old type without breaking existing code.
|
||||
type TensorInternalData = ValueInternalData
|
||||
+4832
File diff suppressed because it is too large
Load Diff
+2167
File diff suppressed because it is too large
Load Diff
+731
@@ -0,0 +1,731 @@
|
||||
// Copyright (c) Microsoft Corporation. All rights reserved.
|
||||
// Licensed under the MIT License.
|
||||
|
||||
// This file contains the training c apis.
|
||||
|
||||
#pragma once
|
||||
#include <stdbool.h>
|
||||
#include "onnxruntime_c_api.h"
|
||||
|
||||
/** \page training_c_cpp_api Training C & C++ APIs
|
||||
*
|
||||
* Training C and C++ APIs are an extension of the \ref c_cpp_api "onnxruntime core C and C++ APIs" and should be used in conjunction with them.
|
||||
*
|
||||
* In order to train a model with onnxruntime, the following training artifacts must be generated:
|
||||
* - The training onnx model
|
||||
* - The checkpoint file
|
||||
* - The optimizer onnx model
|
||||
* - The eval onnx model model (optional)
|
||||
*
|
||||
* These training artifacts can be generated as part of an offline step using the python [utilities](https://github.com/microsoft/onnxruntime/blob/main/orttraining/orttraining/python/training/onnxblock/README.md) made available in the `onnxruntime-training` python package.
|
||||
*
|
||||
* After these artifacts have been generated, the C and C++ utilities listed in this documentation can be leveraged to perform training.
|
||||
*
|
||||
* If any problem is encountered, please create an [issue](https://github.com/microsoft/onnxruntime/issues/new) with your scenario and requirements, and we will be sure to respond and follow up on the request.
|
||||
*
|
||||
* <h1>Training C API</h1>
|
||||
*
|
||||
* ::OrtTrainingApi - Training C API functions.
|
||||
*
|
||||
* This C structure contains functions that enable users to perform training with onnxruntime.
|
||||
*
|
||||
* _Sample Code_:
|
||||
*
|
||||
* ```c
|
||||
* #include <onnxruntime_training_api.h>
|
||||
*
|
||||
* OrtApi* g_ort_api = OrtGetApiBase()->GetApi(ORT_API_VERSION);
|
||||
* OrtTrainingApi* g_ort_training_api = g_ort_api->GetTrainingApi(ORT_API_VERSION);
|
||||
*
|
||||
* OrtEnv* env = NULL;
|
||||
* g_ort_api->CreateEnv(logging_level, logid, &env);
|
||||
* OrtSessionOptions* session_options = NULL;
|
||||
* g_ort_api->CreateSessionOptions(&session_options);
|
||||
*
|
||||
* OrtCheckpointState* state = NULL;
|
||||
* g_ort_training_api->LoadCheckpoint(path_to_checkpoint, &state);
|
||||
*
|
||||
* OrtTrainingSession* training_session = NULL;
|
||||
* g_ort_training_api->CreateTrainingSession(env, session_options, training_model_path,
|
||||
* state, eval_model_path, optimizer_model_path,
|
||||
* &training_session);
|
||||
* // Training loop
|
||||
* {
|
||||
* g_ort_training_api->TrainStep(...);
|
||||
* g_ort_training_api->OptimizerStep(...);
|
||||
* g_ort_training_api->LazyResetGrad(...);
|
||||
* }
|
||||
*
|
||||
* g_ort_training_api->ExportModelForInferencing(training_session, inference_model_path, ...);
|
||||
* g_ort_training_api->SaveCheckpoint(state, path_to_checkpoint, false);
|
||||
*
|
||||
* g_ort_training_api->ReleaseTrainingSession(training_session);
|
||||
* g_ort_training_api->ReleaseCheckpointState(state);
|
||||
* ```
|
||||
*
|
||||
* > **Note**
|
||||
* > The ::OrtCheckpointState contains the entire training state that the ::OrtTrainingSession uses. As a result, the training session must always have access to the state. That is to say, the ::OrtCheckpointState instance must outlive the lifetime of the ::OrtTrainingSession instance.
|
||||
*
|
||||
* <h1>Training C++ API</h1>
|
||||
*
|
||||
* @ref TrainingCpp - Training C++ API classes and functions.
|
||||
*
|
||||
* These C++ classes and functions enable users to perform training with onnxruntime.
|
||||
*
|
||||
* _Sample Code_:
|
||||
*
|
||||
* ```cc
|
||||
* #include <onnxruntime_training_cxx_api.h>
|
||||
*
|
||||
* Ort::Env env;
|
||||
* Ort::SessionOptions session_options;
|
||||
*
|
||||
* auto state = Ort::CheckpointState::LoadCheckpoint(path_to_checkpoint);
|
||||
* auto training_session = Ort::TrainingSession(env, session_options, state, training_model_path,
|
||||
* eval_model_path, optimizer_model_path);
|
||||
*
|
||||
* // Training Loop
|
||||
* {
|
||||
* training_session.TrainStep(...);
|
||||
* training_session.OptimizerStep(...);
|
||||
* training_session.LazyResetGrad(...);
|
||||
* }
|
||||
*
|
||||
* training_session->ExportModelForInferencing(inference_model_path, ...);
|
||||
* Ort::CheckpointState::SaveCheckpoint(state, path_to_checkpoint, false);
|
||||
* ```
|
||||
* > **Note**
|
||||
* > The ::Ort::CheckpointState contains the entire training state that the ::Ort::TrainingSession uses. As a result, the training session must always have access to the state. That is to say, the ::Ort::CheckpointState instance must outlive the lifetime of the ::Ort::TrainingSession instance.
|
||||
*/
|
||||
|
||||
/** @defgroup TrainingC Ort Training C API
|
||||
* @{
|
||||
*/
|
||||
ORT_RUNTIME_CLASS(TrainingSession); // Type that enables performing training for the given user models.
|
||||
ORT_RUNTIME_CLASS(CheckpointState); // Type that holds the training states for the training session.
|
||||
|
||||
/** \brief Type of property to be added to or returned from the ::OrtCheckpointState.
|
||||
*/
|
||||
typedef enum OrtPropertyType {
|
||||
OrtIntProperty = 0,
|
||||
OrtFloatProperty = 1,
|
||||
OrtStringProperty = 2,
|
||||
} OrtPropertyType;
|
||||
|
||||
/** \brief The Training C API that holds onnxruntime training function pointers
|
||||
*
|
||||
* All the Training C API functions are defined inside this structure as pointers to functions.
|
||||
* Call OrtApi::GetTrainingApi to get a pointer to this struct.
|
||||
*
|
||||
* \nosubgrouping
|
||||
*/
|
||||
struct OrtTrainingApi {
|
||||
/// \name Accessing The Training Session State
|
||||
/// @{
|
||||
|
||||
/** \brief Load a checkpoint state from a file on disk into checkpoint_state.
|
||||
*
|
||||
* This function will parse a checkpoint file, pull relevant data and load the training
|
||||
* state into the checkpoint_state. This checkpoint state can then be used to create the
|
||||
* training session by invoking OrtTrainingApi::CreateTrainingSession. By doing so, the training
|
||||
* session will resume training from the given checkpoint state.
|
||||
* \note Note that the training session created with a checkpoint state uses this state to store the entire
|
||||
* training state (including model parameters, its gradients, the optimizer states and the properties).
|
||||
* As a result, it is required that the checkpoint state outlive the lifetime of the training session.
|
||||
* \note Note that the checkpoint file can be either the complete checkpoint or the nominal checkpoint.
|
||||
*
|
||||
* \param[in] checkpoint_path Path to the checkpoint file
|
||||
* \param[out] checkpoint_state Checkpoint state that contains the states of the training session.
|
||||
*
|
||||
* \snippet{doc} snippets.dox OrtStatus Return Value
|
||||
*
|
||||
*/
|
||||
ORT_API2_STATUS(LoadCheckpoint, _In_ const ORTCHAR_T* checkpoint_path,
|
||||
_Outptr_ OrtCheckpointState** checkpoint_state);
|
||||
|
||||
/** \brief Save the given state to a checkpoint file on disk.
|
||||
*
|
||||
* This function serializes the provided checkpoint state to a file on disk.
|
||||
* This checkpoint can later be loaded by invoking OrtTrainingApi::LoadCheckpoint to resume
|
||||
* training from this snapshot of the state.
|
||||
*
|
||||
* \param[in] checkpoint_state The checkpoint state to save.
|
||||
* \param[in] checkpoint_path Path to the checkpoint file.
|
||||
* \param[in] include_optimizer_state Flag to indicate whether to save the optimizer state or not.
|
||||
*
|
||||
* \snippet{doc} snippets.dox OrtStatus Return Value
|
||||
*
|
||||
*/
|
||||
ORT_API2_STATUS(SaveCheckpoint, _In_ OrtCheckpointState* checkpoint_state, _In_ const ORTCHAR_T* checkpoint_path,
|
||||
const bool include_optimizer_state);
|
||||
|
||||
/// @}
|
||||
|
||||
/// \name Implementing The Training Loop
|
||||
/// @{
|
||||
/** \brief Create a training session that can be used to begin or resume training.
|
||||
*
|
||||
* This function creates a training session based on the env and session options provided that can
|
||||
* begin or resume training from a given checkpoint state for the given onnx models.
|
||||
* The checkpoint state represents the parameters of the training session which will be moved
|
||||
* to the device specified by the user through the session options (if necessary).
|
||||
* The training session requires four training artifacts
|
||||
* - The training onnx model
|
||||
* - The evaluation onnx model (optional)
|
||||
* - The optimizer onnx model
|
||||
* - The checkpoint file
|
||||
*
|
||||
* These artifacts can be generated using the `onnxruntime-training` python [utility](https://github.com/microsoft/onnxruntime/blob/main/orttraining/orttraining/python/training/onnxblock/README.md).
|
||||
*
|
||||
* \param[in] env Environment to be used for the training session.
|
||||
* \param[in] options Session options that the user can customize for this training session.
|
||||
* \param[in] checkpoint_state Training states that the training session uses as a starting point for training.
|
||||
* \param[in] train_model_path Model to be used to perform training.
|
||||
* \param[in] eval_model_path Model to be used to perform evaluation.
|
||||
* \param[in] optimizer_model_path Model to be used to perform gradient descent.
|
||||
* \param[out] out Created training session.
|
||||
*
|
||||
* \snippet{doc} snippets.dox OrtStatus Return Value
|
||||
*
|
||||
*/
|
||||
ORT_API2_STATUS(CreateTrainingSession, _In_ const OrtEnv* env, _In_ const OrtSessionOptions* options,
|
||||
_Inout_ OrtCheckpointState* checkpoint_state, _In_ const ORTCHAR_T* train_model_path,
|
||||
_In_ const ORTCHAR_T* eval_model_path, _In_ const ORTCHAR_T* optimizer_model_path,
|
||||
_Outptr_result_maybenull_ OrtTrainingSession** out);
|
||||
|
||||
/** \brief Create a training session that can be used to begin or resume training.
|
||||
* This api provides a way to load all the training artifacts from buffers instead of files.
|
||||
*
|
||||
* \param[in] env Environment to be used for the training session.
|
||||
* \param[in] options Session options that the user can customize for this training session.
|
||||
* \param[in] checkpoint_state Training states that the training session uses as a starting point for training.
|
||||
* \param[in] train_model_data Buffer containing the model data to be used to perform training
|
||||
* \param[in] train_data_length Length of the buffer containing train_model_data
|
||||
* \param[in] eval_model_data Buffer containing the model data to be used to perform evaluation
|
||||
* \param[in] eval_data_length Length of the buffer containing eval_model_data
|
||||
* \param[in] optim_model_data Buffer containing the model data to be used to perform weight update
|
||||
* \param[in] optim_data_length Length of the buffer containing optim_model_data
|
||||
* \param[out] out Created training session.
|
||||
*
|
||||
*/
|
||||
ORT_API2_STATUS(CreateTrainingSessionFromBuffer, _In_ const OrtEnv* env,
|
||||
_In_ const OrtSessionOptions* options, _Inout_ OrtCheckpointState* checkpoint_state,
|
||||
_In_ const void* train_model_data, size_t train_data_length,
|
||||
_In_ const void* eval_model_data, size_t eval_data_length,
|
||||
_In_ const void* optim_model_data, size_t optim_data_length,
|
||||
_Outptr_result_maybenull_ OrtTrainingSession** out);
|
||||
|
||||
/// @}
|
||||
|
||||
/// \name Model IO Information
|
||||
/// @{
|
||||
|
||||
/** \brief Retrieves the number of user outputs in the training model.
|
||||
*
|
||||
* This function returns the number of outputs of the training model so that the user can
|
||||
* allocate space for the number of outputs when OrtTrainingApi::TrainStep is invoked.
|
||||
*
|
||||
* \param[in] sess The `this` pointer to the training session.
|
||||
* \param[out] out Number of user outputs in the training model.
|
||||
*
|
||||
* \snippet{doc} snippets.dox OrtStatus Return Value
|
||||
*
|
||||
*/
|
||||
ORT_API2_STATUS(TrainingSessionGetTrainingModelOutputCount, _In_ const OrtTrainingSession* sess, _Out_ size_t* out);
|
||||
|
||||
/** \brief Retrieves the number of user outputs in the eval model.
|
||||
*
|
||||
* This function returns the number of outputs of the eval model so that the user can
|
||||
* allocate space for the number of outputs when OrtTrainingApi::EvalStep is invoked.
|
||||
*
|
||||
* \param[in] sess The `this` pointer to the training session.
|
||||
* \param[out] out Number of user outputs in the eval model.
|
||||
*
|
||||
* \snippet{doc} snippets.dox OrtStatus Return Value
|
||||
*
|
||||
*/
|
||||
ORT_API2_STATUS(TrainingSessionGetEvalModelOutputCount, _In_ const OrtTrainingSession* sess, _Out_ size_t* out);
|
||||
|
||||
/** \brief Retrieves the names of user outputs in the training model.
|
||||
*
|
||||
* This function returns the names of outputs of the training model that can be associated with the OrtValue(s)
|
||||
* returned by the OrtTrainingApi::TrainStep function.
|
||||
*
|
||||
* \param[in] sess The `this` pointer to the training session.
|
||||
* \param[in] index Index of the output name requested.
|
||||
* \param[in] allocator Allocator to use to allocate the memory for the name.
|
||||
* \param[out] output Name of the training model output at the given index.
|
||||
*
|
||||
* \snippet{doc} snippets.dox OrtStatus Return Value
|
||||
*
|
||||
*/
|
||||
ORT_API2_STATUS(TrainingSessionGetTrainingModelOutputName, _In_ const OrtTrainingSession* sess, size_t index, _Inout_ OrtAllocator* allocator, _Outptr_ char** output);
|
||||
|
||||
/** \brief Retrieves the names of user outputs in the eval model.
|
||||
*
|
||||
* This function returns the names of outputs of the eval model that can be associated with the OrtValue(s) returned
|
||||
* by the OrtTrainingApi::EvalStep function.
|
||||
*
|
||||
* \param[in] sess The `this` pointer to the training session.
|
||||
* \param[in] index Index of the output name requested.
|
||||
* \param[in] allocator Allocator to use to allocate the memory for the name.
|
||||
* \param[out] output Name of the eval model output at the given index.
|
||||
*
|
||||
* \snippet{doc} snippets.dox OrtStatus Return Value
|
||||
*
|
||||
*/
|
||||
ORT_API2_STATUS(TrainingSessionGetEvalModelOutputName, _In_ const OrtTrainingSession* sess, size_t index, _Inout_ OrtAllocator* allocator, _Outptr_ char** output);
|
||||
|
||||
/// @}
|
||||
|
||||
/// \name Implementing The Training Loop
|
||||
/// @{
|
||||
|
||||
/** \brief Reset the gradients of all trainable parameters to zero lazily.
|
||||
*
|
||||
* This function sets the internal state of the training session such that the gradients of the trainable
|
||||
* parameters in the OrtCheckpointState will be scheduled to be reset just before the new gradients are
|
||||
* computed on the next invocation of the next OrtTrainingApi::TrainStep.
|
||||
*
|
||||
* \param[in] session The `this` pointer to the training session.
|
||||
*
|
||||
* \snippet{doc} snippets.dox OrtStatus Return Value
|
||||
*
|
||||
*/
|
||||
ORT_API2_STATUS(LazyResetGrad, _Inout_ OrtTrainingSession* session);
|
||||
|
||||
/** \brief Computes the outputs of the training model and the gradients of the trainable parameters for the given inputs
|
||||
*
|
||||
* This function performs a training step that computes the outputs of the training model and the gradients
|
||||
* of the trainable parameters for the given inputs. The train step is performed based on the training model
|
||||
* that was provided to the training session.
|
||||
* The OrtTrainingApi::TrainStep is equivalent of running forward propagation and backward propagation in a single
|
||||
* step.
|
||||
* The gradients computed are stored inside the training session state so they can be later consumed
|
||||
* by the OrtTrainingApi::OptimizerStep function.
|
||||
* The gradients can be lazily reset by invoking the OrtTrainingApi::LazyResetGrad function.
|
||||
*
|
||||
* \param[in] sess The `this` pointer to the training session.
|
||||
* \param[in] run_options Run options for this training step.
|
||||
* \param[in] inputs_len Number of user inputs to the training model.
|
||||
* \param[in] inputs The user inputs to the training model.
|
||||
* \param[in] outputs_len Number of user outputs expected from this training step.
|
||||
* \param[out] outputs User outputs computed by train step.
|
||||
*
|
||||
* \snippet{doc} snippets.dox OrtStatus Return Value
|
||||
*
|
||||
*/
|
||||
ORT_API2_STATUS(TrainStep, _Inout_ OrtTrainingSession* sess, _In_opt_ const OrtRunOptions* run_options,
|
||||
_In_ size_t inputs_len, _In_reads_(inputs_len) const OrtValue* const* inputs,
|
||||
_In_ size_t outputs_len, _Inout_updates_all_(outputs_len) OrtValue** outputs);
|
||||
|
||||
/** \brief Computes the outputs for the eval model for the given inputs
|
||||
*
|
||||
* This function performs an eval step that computes the outputs of the eval model for the given inputs.
|
||||
* The eval step is performed based on the eval model that was provided to the training session.
|
||||
*
|
||||
* \param[in] sess The `this` pointer to the training session.
|
||||
* \param[in] run_options Run options for this eval step.
|
||||
* \param[in] inputs_len Number of user inputs to the eval model.
|
||||
* \param[in] inputs The user inputs to the eval model.
|
||||
* \param[in] outputs_len Number of user outputs expected from this eval step.
|
||||
* \param[out] outputs User outputs computed by eval step.
|
||||
*
|
||||
* \snippet{doc} snippets.dox OrtStatus Return Value
|
||||
*
|
||||
*/
|
||||
ORT_API2_STATUS(EvalStep, _In_ const OrtTrainingSession* sess, _In_opt_ const OrtRunOptions* run_options,
|
||||
_In_ size_t inputs_len, _In_reads_(inputs_len) const OrtValue* const* inputs,
|
||||
_In_ size_t outputs_len, _Inout_updates_all_(outputs_len) OrtValue** outputs);
|
||||
|
||||
/** \brief Sets the learning rate for this training session.
|
||||
*
|
||||
* This function allows users to set the learning rate for the training session. The current
|
||||
* learning rate is maintained by the training session and can be overwritten by invoking
|
||||
* this function with the desired learning rate. This function should not be used when a valid
|
||||
* learning rate scheduler is registered. It should be used either to set the learning rate
|
||||
* derived from a custom learning rate scheduler or to set a constant learning rate to be used
|
||||
* throughout the training session.
|
||||
* \note Please note that this function does not set the initial learning rate that may be needed
|
||||
* by the predefined learning rate schedulers. To set the initial learning rate for learning
|
||||
* rate schedulers, please look at the function OrtTrainingApi::RegisterLinearLRScheduler.
|
||||
*
|
||||
* \param[in] sess The `this` pointer to the training session.
|
||||
* \param[in] learning_rate Desired learning rate to be set.
|
||||
*
|
||||
* \snippet{doc} snippets.dox OrtStatus Return Value
|
||||
*
|
||||
*/
|
||||
ORT_API2_STATUS(SetLearningRate, _Inout_ OrtTrainingSession* sess, _In_ float learning_rate);
|
||||
|
||||
/** \brief Gets the current learning rate for this training session.
|
||||
*
|
||||
* This function allows users to get the learning rate for the training session. The current
|
||||
* learning rate is maintained by the training session, and users can query it for the purpose
|
||||
* of implementing their own learning rate schedulers.
|
||||
*
|
||||
* \param[in] sess The `this` pointer to the training session.
|
||||
* \param[out] learning_rate Learning rate currently in use by the training session.
|
||||
*
|
||||
* \snippet{doc} snippets.dox OrtStatus Return Value
|
||||
*
|
||||
*/
|
||||
ORT_API2_STATUS(GetLearningRate, _Inout_ OrtTrainingSession* sess, _Out_ float* learning_rate);
|
||||
|
||||
/** \brief Performs the weight updates for the trainable parameters using the optimizer model.
|
||||
*
|
||||
* This function performs the weight update step that updates the trainable parameters such that they
|
||||
* take a step in the direction of their gradients (gradient descent). The optimizer step is performed
|
||||
* based on the optimizer model that was provided to the training session.
|
||||
* The updated parameters are stored inside the training state so that they can be used by the next
|
||||
* OrtTrainingApi::TrainStep function call.
|
||||
*
|
||||
* \param[in] sess The `this` pointer to the training session.
|
||||
* \param[in] run_options Run options for this optimizer step.
|
||||
*
|
||||
* \snippet{doc} snippets.dox OrtStatus Return Value
|
||||
*
|
||||
*/
|
||||
ORT_API2_STATUS(OptimizerStep, _Inout_ OrtTrainingSession* sess,
|
||||
_In_opt_ const OrtRunOptions* run_options);
|
||||
|
||||
/** \brief Registers a linear learning rate scheduler for the training session.
|
||||
*
|
||||
* Register a linear learning rate scheduler that decays the learning rate by linearly updated
|
||||
* multiplicative factor from the initial learning rate set on the training session to 0. The decay
|
||||
* is performed after the initial warm up phase where the learning rate is linearly incremented
|
||||
* from 0 to the initial learning rate provided.
|
||||
*
|
||||
* \param[in] sess The `this` pointer to the training session.
|
||||
* \param[in] warmup_step_count Warmup steps for LR warmup.
|
||||
* \param[in] total_step_count Total step count.
|
||||
* \param[in] initial_lr The initial learning rate to be used by the training session.
|
||||
*
|
||||
* \snippet{doc} snippets.dox OrtStatus Return Value
|
||||
*
|
||||
*/
|
||||
ORT_API2_STATUS(RegisterLinearLRScheduler, _Inout_ OrtTrainingSession* sess, _In_ const int64_t warmup_step_count,
|
||||
_In_ const int64_t total_step_count, _In_ const float initial_lr);
|
||||
|
||||
/** \brief Update the learning rate based on the registered learing rate scheduler.
|
||||
*
|
||||
* Takes a scheduler step that updates the learning rate that is being used by the training session.
|
||||
* This function should typically be called before invoking the optimizer step for each round,
|
||||
* or as determined necessary to update the learning rate being used by the training session.
|
||||
* \note Please note that a valid predefined learning rate scheduler must be first registered to invoke this
|
||||
* function.
|
||||
*
|
||||
* \param[in] sess The `this` pointer to the training session.
|
||||
*
|
||||
* \snippet{doc} snippets.dox OrtStatus Return Value
|
||||
*
|
||||
*/
|
||||
ORT_API2_STATUS(SchedulerStep, _Inout_ OrtTrainingSession* sess);
|
||||
|
||||
/// @}
|
||||
|
||||
/// \name Accessing The Training Session State
|
||||
/// @{
|
||||
/** \brief Retrieves the size of all the parameters.
|
||||
*
|
||||
* Calculates the total number of primitive (datatype of the parameters) elements of all the parameters in the
|
||||
* training state.
|
||||
* When trainable_only argument is true, the size is calculated for trainable params only.
|
||||
*
|
||||
* \param[in] sess The `this` pointer to the training session.
|
||||
* \param[out] out Size of all parameter elements.
|
||||
* \param[in] trainable_only Whether to skip non-trainable parameters
|
||||
*
|
||||
* \snippet{doc} snippets.dox OrtStatus Return Value
|
||||
*
|
||||
*/
|
||||
ORT_API2_STATUS(GetParametersSize, _Inout_ OrtTrainingSession* sess, _Out_ size_t* out, bool trainable_only);
|
||||
|
||||
/** \brief Copy all parameters to a contiguous buffer held by the argument parameters_buffer
|
||||
*
|
||||
* The parameters_buffer has to be of the size given by GetParametersSize api call,
|
||||
* with matching setting for the argument trainable_only. All the target parameters must be of the same
|
||||
* datatype. The OrtValue must be pre-allocated onto
|
||||
* the desired device. This is a complementary function to OrtTrainingApi::CopyBufferToParameters.
|
||||
* Parameter ordering is preserved.
|
||||
* User is responsible for allocating and freeing the resources used by the parameters_buffer.
|
||||
*
|
||||
* \param[in] sess The `this` pointer to the training session.
|
||||
* \param[in] trainable_only Whether to skip non-trainable parameters
|
||||
* \param[out] parameters_buffer The pre-allocated OrtValue buffer to copy onto.
|
||||
*
|
||||
* \snippet{doc} snippets.dox OrtStatus Return Value
|
||||
*
|
||||
*/
|
||||
ORT_API2_STATUS(CopyParametersToBuffer, _Inout_ OrtTrainingSession* sess,
|
||||
_Inout_ OrtValue* parameters_buffer, bool trainable_only);
|
||||
|
||||
/** \brief Copy parameter values from the given contiguous buffer held by parameters_buffer to the training state
|
||||
*
|
||||
* The parameters_buffer argument has to be of the size given by OrtTrainingApi::GetParametersSize api call,
|
||||
* with matching setting for trainable_only argument. All the target parameters must be of the same
|
||||
* datatype. This is a complementary function to OrtTrainingApi::CopyParametersToBuffer
|
||||
* and can be used to load updated buffer values onto the training state.
|
||||
* Parameter ordering is preserved.
|
||||
* User is responsible for allocating and freeing the resources used by the parameters_buffer.
|
||||
* In case the training session was created with a nominal checkpoint, invoking this function is required
|
||||
* to load the updated parameters onto the checkpoint to complete it.
|
||||
*
|
||||
* \param[in] sess The `this` pointer to the training session.
|
||||
* \param[in] trainable_only Whether to skip non-trainable parameters
|
||||
* \param[out] parameters_buffer The pre-allocated OrtValue buffer to copy from.
|
||||
*
|
||||
* \snippet{doc} snippets.dox OrtStatus Return Value
|
||||
*
|
||||
*/
|
||||
ORT_API2_STATUS(CopyBufferToParameters, _Inout_ OrtTrainingSession* sess,
|
||||
_Inout_ OrtValue* parameters_buffer, bool trainable_only);
|
||||
|
||||
/// @}
|
||||
|
||||
/// \name Release Training Resources
|
||||
/// @{
|
||||
|
||||
/** \brief Frees up the memory used up by the training session.
|
||||
*
|
||||
* This function frees up any memory that was allocated in the training session. The training
|
||||
* session can no longer be used after this call.
|
||||
*
|
||||
*/
|
||||
ORT_CLASS_RELEASE(TrainingSession);
|
||||
|
||||
/** \brief Frees up the memory used up by the checkpoint state.
|
||||
*
|
||||
* This function frees up any memory that was allocated in the checkpoint state. The checkpoint
|
||||
* state can no longer be used after this call.
|
||||
* \note Note that the checkpoint state must be released only after the training session has been released.
|
||||
*
|
||||
*/
|
||||
ORT_CLASS_RELEASE(CheckpointState);
|
||||
|
||||
/// @}
|
||||
|
||||
/// \name Prepare For Inferencing
|
||||
/// @{
|
||||
/** \brief Export a model that can be used for inferencing.
|
||||
*
|
||||
* If the training session was provided with an eval model, the training session can generate
|
||||
* an inference model if it knows the inference graph outputs. The input inference graph outputs
|
||||
* are used to prune the eval model so that the inference model's outputs align with the provided outputs.
|
||||
* The exported model is saved at the path provided and can be used for inferencing with InferenceSession.
|
||||
* \note Note that the function re-loads the eval model from the path provided to OrtTrainingApi::CreateTrainingSession
|
||||
* and expects that this path still be valid.
|
||||
*
|
||||
* \param[in] sess The `this` pointer to the training session.
|
||||
* \param[in] inference_model_path Path where the inference model should be serialized to.
|
||||
* \param[in] graph_outputs_len Size of the graph output names array.
|
||||
* \param[in] graph_output_names Names of the outputs that are needed in the inference model.
|
||||
*
|
||||
* \snippet{doc} snippets.dox OrtStatus Return Value
|
||||
*
|
||||
*/
|
||||
ORT_API2_STATUS(ExportModelForInferencing, _Inout_ OrtTrainingSession* sess,
|
||||
_In_ const ORTCHAR_T* inference_model_path, size_t graph_outputs_len,
|
||||
_In_reads_(graph_outputs_len) const char* const* graph_output_names);
|
||||
|
||||
/// @}
|
||||
|
||||
/// \name Training Utilities
|
||||
/// @{
|
||||
/** \brief Sets the seed used for random number generation in Onnxruntime.
|
||||
*
|
||||
* Use this function to generate reproducible results. It should be noted that completely reproducible
|
||||
* results are not guaranteed.
|
||||
*
|
||||
* \param[in] seed The seed to be set.
|
||||
*
|
||||
* \snippet{doc} snippets.dox OrtStatus Return Value
|
||||
*
|
||||
*/
|
||||
ORT_API2_STATUS(SetSeed, _In_ const int64_t seed);
|
||||
|
||||
/// @}
|
||||
|
||||
/// \name Model IO Information
|
||||
/// @{
|
||||
/** \brief Retrieves the number of user inputs in the training model.
|
||||
*
|
||||
* This function returns the number of inputs of the training model so that the user can accordingly
|
||||
* allocate the OrtValue(s) provided to the OrtTrainingApi::TrainStep function.
|
||||
*
|
||||
* \param[in] sess The `this` pointer to the training session.
|
||||
* \param[out] out Number of user inputs in the training model.
|
||||
*
|
||||
* \snippet{doc} snippets.dox OrtStatus Return Value
|
||||
*
|
||||
*/
|
||||
ORT_API2_STATUS(TrainingSessionGetTrainingModelInputCount, _In_ const OrtTrainingSession* sess, _Out_ size_t* out);
|
||||
|
||||
/** \brief Retrieves the number of user inputs in the eval model.
|
||||
*
|
||||
* This function returns the number of inputs of the eval model so that the user can accordingly
|
||||
* allocate the OrtValue(s) provided to the OrtTrainingApi::EvalStep function.
|
||||
*
|
||||
* \param[in] sess The `this` pointer to the training session.
|
||||
* \param[out] out Number of user inputs in the eval model.
|
||||
*
|
||||
* \snippet{doc} snippets.dox OrtStatus Return Value
|
||||
*
|
||||
*/
|
||||
ORT_API2_STATUS(TrainingSessionGetEvalModelInputCount, _In_ const OrtTrainingSession* sess, _Out_ size_t* out);
|
||||
|
||||
/** \brief Retrieves the name of the user input at given index in the training model.
|
||||
*
|
||||
* This function returns the names of inputs of the training model that can be associated with the
|
||||
* OrtValue(s) provided to the OrtTrainingApi::TrainStep function.
|
||||
*
|
||||
* \param[in] sess The `this` pointer to the training session.
|
||||
* \param[in] index The index of the training model input name requested.
|
||||
* \param[in] allocator The allocator to use to allocate the memory for the requested name.
|
||||
* \param[out] output Name of the user input for the training model at the given index.
|
||||
*
|
||||
* \snippet{doc} snippets.dox OrtStatus Return Value
|
||||
*
|
||||
*/
|
||||
ORT_API2_STATUS(TrainingSessionGetTrainingModelInputName, _In_ const OrtTrainingSession* sess, size_t index,
|
||||
_In_ OrtAllocator* allocator, _Outptr_ char** output);
|
||||
|
||||
/** \brief Retrieves the name of the user input at given index in the eval model.
|
||||
*
|
||||
* This function returns the names of inputs of the eval model that can be associated with the OrtValue(s) provided
|
||||
* to the OrtTrainingApi::EvalStep function.
|
||||
*
|
||||
* \param[in] sess The `this` pointer to the training session.
|
||||
* \param[in] index The index of the eval model input name requested.
|
||||
* \param[in] allocator The allocator to use to allocate the memory for the requested name.
|
||||
* \param[out] output Name of the user input for the eval model at the given index.
|
||||
*
|
||||
* \snippet{doc} snippets.dox OrtStatus Return Value
|
||||
*
|
||||
*/
|
||||
ORT_API2_STATUS(TrainingSessionGetEvalModelInputName, _In_ const OrtTrainingSession* sess, size_t index,
|
||||
_In_ OrtAllocator* allocator, _Outptr_ char** output);
|
||||
|
||||
/// @}
|
||||
|
||||
/// \name Accessing The Training Session State
|
||||
/// @{
|
||||
|
||||
/** \brief Adds or updates the given property to/in the checkpoint state.
|
||||
*
|
||||
* Runtime properties such as epoch, training step, best score, and others can be added to the checkpoint
|
||||
* state by the user by calling this function with the corresponding property name and value.
|
||||
* The given property name must be unique to be able to successfully add the property.
|
||||
*
|
||||
* \param[in] checkpoint_state The checkpoint state which should hold the property.
|
||||
* \param[in] property_name Name of the property being added or updated.
|
||||
* \param[in] property_type Type of the property associated with the given name.
|
||||
* \param[in] property_value Property value associated with the given name.
|
||||
*
|
||||
* \snippet{doc} snippets.dox OrtStatus Return Value
|
||||
*
|
||||
*/
|
||||
ORT_API2_STATUS(AddProperty, _Inout_ OrtCheckpointState* checkpoint_state,
|
||||
_In_ const char* property_name, _In_ enum OrtPropertyType property_type,
|
||||
_In_ void* property_value);
|
||||
|
||||
/** \brief Gets the property value associated with the given name from the checkpoint state.
|
||||
*
|
||||
* Gets the property value from an existing entry in the checkpoint state. The property must
|
||||
* exist in the checkpoint state to be able to retrieve it successfully.
|
||||
*
|
||||
* \param[in] checkpoint_state The checkpoint state that is currently holding the property.
|
||||
* \param[in] property_name Name of the property being retrieved.
|
||||
* \param[in] allocator Allocator used to allocate the memory for the property_value.
|
||||
* \param[out] property_type Type of the property associated with the given name.
|
||||
* \param[out] property_value Property value associated with the given name.
|
||||
*
|
||||
* \snippet{doc} snippets.dox OrtStatus Return Value
|
||||
*
|
||||
*/
|
||||
ORT_API2_STATUS(GetProperty, _In_ const OrtCheckpointState* checkpoint_state,
|
||||
_In_ const char* property_name, _Inout_ OrtAllocator* allocator,
|
||||
_Out_ enum OrtPropertyType* property_type, _Outptr_ void** property_value);
|
||||
|
||||
/// @}
|
||||
|
||||
/// \name Accessing The Training Session State
|
||||
/// @{
|
||||
|
||||
/** \brief Load a checkpoint state from a buffer into checkpoint_state.
|
||||
*
|
||||
* This function will parse a checkpoint bytes buffer, pull relevant data and load the training
|
||||
* state into the checkpoint_state. This checkpoint state can then be used to create the
|
||||
* training session by invoking OrtTrainingApi::CreateTrainingSession. By doing so, the training
|
||||
* session will resume training from the given checkpoint state.
|
||||
* \note Note that the training session created with a checkpoint state uses this state to store the entire
|
||||
* training state (including model parameters, its gradients, the optimizer states and the properties).
|
||||
* As a result, it is required that the checkpoint state outlive the lifetime of the training session.
|
||||
*
|
||||
* \param[in] checkpoint_buffer Path to the checkpoint bytes buffer.
|
||||
* \param[in] num_bytes Number of bytes in the checkpoint buffer.
|
||||
* \param[out] checkpoint_state Checkpoint state that contains the states of the training session.
|
||||
*
|
||||
* \snippet{doc} snippets.dox OrtStatus Return Value
|
||||
*
|
||||
*/
|
||||
ORT_API2_STATUS(LoadCheckpointFromBuffer, _In_ const void* checkpoint_buffer,
|
||||
_In_ const size_t num_bytes, _Outptr_ OrtCheckpointState** checkpoint_state);
|
||||
|
||||
/** \brief Retrieves the type and shape information of the parameter associated with the given parameter name.
|
||||
*
|
||||
* This function retrieves the type and shape of the parameter associated with the given parameter name.
|
||||
* The parameter must exist in the checkpoint state to be able to retrieve its type and shape information successfully.
|
||||
*
|
||||
* \param[in] checkpoint_state The checkpoint state.
|
||||
* \param[in] parameter_name Name of the parameter being retrieved.
|
||||
* \param[out] parameter_type_and_shape The type and shape of the parameter being retrieved.
|
||||
*
|
||||
* \snippet{doc} snippets.dox OrtStatus Return Value
|
||||
*
|
||||
*/
|
||||
ORT_API2_STATUS(GetParameterTypeAndShape, _In_ const OrtCheckpointState* checkpoint_state,
|
||||
_In_ const char* parameter_name, _Outptr_ OrtTensorTypeAndShapeInfo** parameter_type_and_shape);
|
||||
|
||||
/** \brief Updates the data associated with the model parameter in the checkpoint state for the given parameter name.
|
||||
*
|
||||
* This function updates a model parameter in the checkpoint state with the given parameter data.
|
||||
* The training session must be already created with the checkpoint state that contains the parameter
|
||||
* being updated. The given parameter is copied over to the registered device for the training session.
|
||||
* The parameter must exist in the checkpoint state to be able to update it successfully.
|
||||
*
|
||||
* \param[in] checkpoint_state The checkpoint state.
|
||||
* \param[in] parameter_name Name of the parameter being updated.
|
||||
* \param[in] parameter The parameter data that should replace the existing parameter data.
|
||||
*
|
||||
* \snippet{doc} snippets.dox OrtStatus Return Value
|
||||
*
|
||||
*/
|
||||
ORT_API2_STATUS(UpdateParameter, _Inout_ OrtCheckpointState* checkpoint_state,
|
||||
_In_ const char* parameter_name, _In_ OrtValue* parameter);
|
||||
|
||||
/** \brief Gets the data associated with the model parameter from the checkpoint state for the given parameter name.
|
||||
*
|
||||
* This function retrieves the model parameter data from the checkpoint state for the given parameter name.
|
||||
* The parameter is copied over and returned as an OrtValue. The training session must be already created
|
||||
* with the checkpoint state that contains the parameter being retrieved.
|
||||
* The parameter must exist in the checkpoint state to be able to retrieve it successfully.
|
||||
*
|
||||
* \param[in] checkpoint_state The checkpoint state.
|
||||
* \param[in] parameter_name Name of the parameter being retrieved.
|
||||
* \param[in] allocator Allocator used to allocate the memory for the parameter.
|
||||
* \param[out] parameter The parameter data that is retrieved from the checkpoint state.
|
||||
*
|
||||
* \snippet{doc} snippets.dox OrtStatus Return Value
|
||||
*
|
||||
*/
|
||||
ORT_API2_STATUS(GetParameter, _In_ const OrtCheckpointState* checkpoint_state,
|
||||
_In_ const char* parameter_name, _Inout_ OrtAllocator* allocator,
|
||||
_Outptr_ OrtValue** parameter);
|
||||
|
||||
/// @}
|
||||
};
|
||||
|
||||
typedef struct OrtTrainingApi OrtTrainingApi;
|
||||
|
||||
/// @}
|
||||
+639
@@ -0,0 +1,639 @@
|
||||
package onnxruntime_go
|
||||
|
||||
// #cgo CFLAGS: -O2 -g
|
||||
//
|
||||
// #include "onnxruntime_wrapper.h"
|
||||
import "C"
|
||||
import (
|
||||
"fmt"
|
||||
"os"
|
||||
"path/filepath"
|
||||
"unsafe"
|
||||
)
|
||||
|
||||
var trainingNotSupportedError error = fmt.Errorf("training not supported by onnx library")
|
||||
|
||||
// Scalar is like a tensor but the underlying go slice is of length 1 and it has no dimension.
|
||||
// It can be used to store e.g. the loss from a training cycle.
|
||||
type Scalar[T TensorData] struct {
|
||||
data []T
|
||||
dataSize uintptr
|
||||
ortValue *C.OrtValue
|
||||
}
|
||||
|
||||
func (s *Scalar[T]) GetShape() Shape {
|
||||
return nil
|
||||
}
|
||||
|
||||
func (s *Scalar[T]) ZeroContents() {
|
||||
C.memset(unsafe.Pointer(&s.data[0]), 0, C.size_t(s.dataSize))
|
||||
}
|
||||
|
||||
func (s *Scalar[T]) Destroy() error {
|
||||
C.ReleaseOrtValue(s.ortValue)
|
||||
s.ortValue = nil
|
||||
s.data = nil
|
||||
s.dataSize = 0
|
||||
return nil
|
||||
}
|
||||
|
||||
// GetData returns the undelying data for the scalar.
|
||||
// If you want to explicitly set the scalar's data, use Set.
|
||||
func (t *Scalar[T]) GetData() T {
|
||||
return t.data[0]
|
||||
}
|
||||
|
||||
// Set allows to explicitly set the underlying value for the scalar.
|
||||
func (t *Scalar[T]) Set(value T) {
|
||||
t.data = []T{value}
|
||||
}
|
||||
|
||||
func (t *Scalar[T]) DataType() C.ONNXTensorElementDataType {
|
||||
return GetTensorElementDataType[T]()
|
||||
}
|
||||
|
||||
func (t *Scalar[_]) GetInternals() *ValueInternalData {
|
||||
return &ValueInternalData{
|
||||
ortValue: t.ortValue,
|
||||
}
|
||||
}
|
||||
|
||||
func (t *Scalar[_]) GetONNXType() ONNXType {
|
||||
return ONNXTypeTensor
|
||||
}
|
||||
|
||||
// NewEmptyScalar creates a new scalar of type T.
|
||||
func NewEmptyScalar[T TensorData]() (*Scalar[T], error) {
|
||||
var data T
|
||||
return NewScalar(data)
|
||||
}
|
||||
|
||||
// NewScalar creates a new scalar of type T backed by a value of type T.
|
||||
// Note that, differently from tensors, this is not a []T but just a value T.
|
||||
func NewScalar[T TensorData](data T) (*Scalar[T], error) {
|
||||
if !IsInitialized() {
|
||||
return nil, NotInitializedError
|
||||
}
|
||||
|
||||
dataSlice := []T{data}
|
||||
var ortValue *C.OrtValue
|
||||
dataType := GetTensorElementDataType[T]()
|
||||
dataSize := unsafe.Sizeof(dataSlice[0]) * uintptr(1)
|
||||
|
||||
status := C.CreateOrtTensorWithShape(unsafe.Pointer(&dataSlice[0]),
|
||||
C.size_t(dataSize), nil, C.int64_t(0), ortMemoryInfo, dataType, &ortValue)
|
||||
if status != nil {
|
||||
return nil, statusToError(status)
|
||||
}
|
||||
toReturn := Scalar[T]{
|
||||
data: dataSlice,
|
||||
dataSize: dataSize,
|
||||
ortValue: ortValue,
|
||||
}
|
||||
return &toReturn, nil
|
||||
}
|
||||
|
||||
// TraininSession is the type that wraps the C training session object.
|
||||
type TrainingSession struct {
|
||||
ortTrainingSession *C.OrtTrainingSession
|
||||
ortCheckpointState *C.OrtCheckpointState
|
||||
inputs []*C.OrtValue
|
||||
outputs []*C.OrtValue
|
||||
trainingModelPath *C.char
|
||||
optimizerModelPath *C.char
|
||||
evalModelPath *C.char
|
||||
}
|
||||
|
||||
// ExportModel is used to export the final trained model to disk. It requires the path for
|
||||
// the exported model as well as the names of the graph nodes to export.
|
||||
// Note that currently the final model can only be exported if the session has been
|
||||
// initialized with NewTrainingSession and the path to the eval model has been provided.
|
||||
func (s *TrainingSession) ExportModel(path string, outputNames []string) error {
|
||||
if s.evalModelPath == nil {
|
||||
return fmt.Errorf("final model can only be exported if the eval model path is " +
|
||||
"provided at session creation time (see NewTrainingSession)")
|
||||
}
|
||||
if path == "" {
|
||||
return fmt.Errorf("path cannot be empty")
|
||||
}
|
||||
dir, _ := filepath.Split(path)
|
||||
if _, err := os.Stat(dir); dir != "" && os.IsNotExist(err) {
|
||||
return fmt.Errorf("directory %s does not exist", dir)
|
||||
}
|
||||
|
||||
cOutputNames := make([]*C.char, len(outputNames))
|
||||
for i, name := range outputNames {
|
||||
cOutputNames[i] = C.CString(name)
|
||||
}
|
||||
cPath, err := createOrtCharString(path)
|
||||
if err != nil {
|
||||
return fmt.Errorf("Error converting export path to C string: %w", err)
|
||||
}
|
||||
outputLength := C.size_t(len(outputNames))
|
||||
defer func() {
|
||||
for i := range cOutputNames {
|
||||
C.free(unsafe.Pointer(cOutputNames[i]))
|
||||
}
|
||||
C.free(unsafe.Pointer(cPath))
|
||||
}()
|
||||
status := C.ExportModel(s.ortTrainingSession, cPath, outputLength, &cOutputNames[0])
|
||||
if status != nil {
|
||||
return statusToError(status)
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// SaveCheckpoint can be used to save the current checkpoint state at the specified path.
|
||||
// This is useful to snapshot the training parameters to continue training later or on
|
||||
// a different machine.
|
||||
func (s *TrainingSession) SaveCheckpoint(path string, saveOptimizerState bool) error {
|
||||
if path == "" {
|
||||
return fmt.Errorf("path cannot be empty")
|
||||
}
|
||||
dir, _ := filepath.Split(path)
|
||||
if _, err := os.Stat(dir); dir != "" && os.IsNotExist(err) {
|
||||
return fmt.Errorf("directory %s does not exist", dir)
|
||||
}
|
||||
|
||||
cPath, err := createOrtCharString(path)
|
||||
if err != nil {
|
||||
return fmt.Errorf("Error converting path to C string: %w", err)
|
||||
}
|
||||
var saveOptimizer int
|
||||
if saveOptimizerState {
|
||||
saveOptimizer = 1
|
||||
}
|
||||
|
||||
defer func() {
|
||||
C.free(unsafe.Pointer(cPath))
|
||||
}()
|
||||
|
||||
status := C.SaveCheckpoint(s.ortCheckpointState, cPath, C.size_t(saveOptimizer))
|
||||
if status != nil {
|
||||
return statusToError(status)
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// Destroy frees all the C memory associated to a training session.
|
||||
func (s *TrainingSession) Destroy() error {
|
||||
if s.ortTrainingSession != nil {
|
||||
C.ReleaseOrtTrainingSession(s.ortTrainingSession)
|
||||
s.ortTrainingSession = nil
|
||||
}
|
||||
// note: checkpoint MUST be released after session
|
||||
if s.ortCheckpointState != nil {
|
||||
C.ReleaseCheckpointState(s.ortCheckpointState)
|
||||
s.ortCheckpointState = nil
|
||||
}
|
||||
C.free(unsafe.Pointer(s.trainingModelPath))
|
||||
s.trainingModelPath = nil
|
||||
C.free(unsafe.Pointer(s.evalModelPath))
|
||||
s.evalModelPath = nil
|
||||
C.free(unsafe.Pointer(s.optimizerModelPath))
|
||||
s.optimizerModelPath = nil
|
||||
s.inputs = nil
|
||||
s.outputs = nil
|
||||
return nil
|
||||
}
|
||||
|
||||
// TrainStep performs the training step.
|
||||
func (s *TrainingSession) TrainStep() error {
|
||||
inputLength := C.size_t(len(s.inputs))
|
||||
outputLength := C.size_t(len(s.outputs))
|
||||
status := C.TrainStep(s.ortTrainingSession, inputLength, &s.inputs[0], outputLength, &s.outputs[0])
|
||||
if status != nil {
|
||||
return fmt.Errorf("error performing training step: %w", statusToError(status))
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// TrainStep performs the optimizer step.
|
||||
func (s *TrainingSession) OptimizerStep() error {
|
||||
status := C.OptimizerStep(s.ortTrainingSession)
|
||||
if status != nil {
|
||||
return fmt.Errorf("error performing optimizer step: %w", statusToError(status))
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// TrainStep performs the LazyResetGrad step.
|
||||
func (s *TrainingSession) LazyResetGrad() error {
|
||||
status := C.LazyResetGrad(s.ortTrainingSession)
|
||||
if status != nil {
|
||||
return fmt.Errorf("error performing lazyResetGrad step: %w", statusToError(status))
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func getInputName(s *C.OrtTrainingSession, i int, model string) (string, error) {
|
||||
var cName *C.char
|
||||
var status *C.OrtStatus
|
||||
switch model {
|
||||
case "train":
|
||||
status = C.TrainingSessionGetTrainingInputName(s, C.size_t(i), &cName)
|
||||
case "eval":
|
||||
status = C.TrainingSessionGetEvalInputName(s, C.size_t(i), &cName)
|
||||
default:
|
||||
return "", fmt.Errorf("%s model not recognized", model)
|
||||
}
|
||||
if status != nil {
|
||||
return "", fmt.Errorf("error getting name: %w", statusToError(status))
|
||||
}
|
||||
|
||||
name, e := convertORTString(cName)
|
||||
if e != nil {
|
||||
return "", fmt.Errorf("error converting C name to Go string: %w", e)
|
||||
}
|
||||
return name, nil
|
||||
}
|
||||
|
||||
func getOutputName(s *C.OrtTrainingSession, i int, model string) (string, error) {
|
||||
var cName *C.char
|
||||
var status *C.OrtStatus
|
||||
switch model {
|
||||
case "train":
|
||||
status = C.TrainingSessionGetTrainingOutputName(s, C.size_t(i), &cName)
|
||||
case "eval":
|
||||
status = C.TrainingSessionGetEvalOutputName(s, C.size_t(i), &cName)
|
||||
default:
|
||||
return "", fmt.Errorf("%s model not recognized", model)
|
||||
}
|
||||
if status != nil {
|
||||
return "", fmt.Errorf("error getting name: %w", statusToError(status))
|
||||
}
|
||||
name, e := convertORTString(cName)
|
||||
if e != nil {
|
||||
return "", fmt.Errorf("error converting C name to Go string: %w", e)
|
||||
}
|
||||
return name, nil
|
||||
}
|
||||
|
||||
type TrainingInputOutputNames struct {
|
||||
TrainingInputNames []string
|
||||
EvalInputNames []string
|
||||
TrainingOutputNames []string
|
||||
EvalOutputNames []string
|
||||
}
|
||||
|
||||
// GetInputOutputNames returns the names of the training inputs and outputs
|
||||
// for the training and validation models. Eval model is optional and can be empty
|
||||
// string.
|
||||
func GetInputOutputNames(checkpointStatePath string,
|
||||
trainingModelPath string,
|
||||
evalModelPath string) (*TrainingInputOutputNames, error) {
|
||||
|
||||
options, e := NewSessionOptions()
|
||||
if e != nil {
|
||||
return nil, fmt.Errorf("failed creating options with error: %v\n", e)
|
||||
}
|
||||
defer options.Destroy()
|
||||
|
||||
checkpointData, e := os.ReadFile(checkpointStatePath)
|
||||
if e != nil {
|
||||
return nil, fmt.Errorf("error reading %s: %w", checkpointStatePath, e)
|
||||
}
|
||||
|
||||
trainingData, e := os.ReadFile(trainingModelPath)
|
||||
if e != nil {
|
||||
return nil, fmt.Errorf("error reading %s: %w", checkpointStatePath, e)
|
||||
}
|
||||
|
||||
var evalData []byte
|
||||
if evalModelPath != "" {
|
||||
evalData, e = os.ReadFile(evalModelPath)
|
||||
if e != nil {
|
||||
return nil, fmt.Errorf("error reading %s: %w", evalModelPath, e)
|
||||
}
|
||||
}
|
||||
|
||||
// create checkpoint C object
|
||||
ortCheckpointState, e := createCCheckpoint(checkpointData)
|
||||
if e != nil {
|
||||
return nil, fmt.Errorf("error creating C checkpointState: %w", e)
|
||||
}
|
||||
|
||||
// create session C object
|
||||
ortTrainingSession, e := createCTrainingSessionWithOnnxData(ortCheckpointState,
|
||||
trainingData, evalData, nil, options)
|
||||
if e != nil {
|
||||
C.ReleaseCheckpointState(ortCheckpointState)
|
||||
return nil, fmt.Errorf("error creating C training session: %w", e)
|
||||
}
|
||||
defer func() {
|
||||
C.ReleaseOrtTrainingSession(ortTrainingSession)
|
||||
C.ReleaseCheckpointState(ortCheckpointState)
|
||||
}()
|
||||
|
||||
var inputCountTraining, inputCountEval C.size_t
|
||||
status := C.TrainingSessionGetInputCount(ortTrainingSession, &inputCountTraining, &inputCountEval)
|
||||
if status != nil {
|
||||
return nil, statusToError(status)
|
||||
}
|
||||
|
||||
var outputCountTraining, outputCountEval C.size_t
|
||||
status = C.TrainingSessionGetOutputCount(ortTrainingSession, &outputCountTraining, &outputCountEval)
|
||||
if status != nil {
|
||||
return nil, statusToError(status)
|
||||
}
|
||||
|
||||
trainInputNames := make([]string, inputCountTraining)
|
||||
trainOutputNames := make([]string, outputCountTraining)
|
||||
|
||||
for i := 0; i < int(inputCountTraining); i++ {
|
||||
name, err := getInputName(ortTrainingSession, i, "train")
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("error retrieving train input name: %w", err)
|
||||
}
|
||||
trainInputNames[i] = name
|
||||
}
|
||||
|
||||
for i := 0; i < int(outputCountTraining); i++ {
|
||||
name, err := getOutputName(ortTrainingSession, i, "train")
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("error retrieving train output name: %w", err)
|
||||
}
|
||||
trainOutputNames[i] = name
|
||||
}
|
||||
|
||||
var evalInputNames []string
|
||||
var evalOutputNames []string
|
||||
|
||||
if len(evalData) > 0 {
|
||||
evalInputNames = make([]string, inputCountEval)
|
||||
evalOutputNames = make([]string, outputCountEval)
|
||||
|
||||
for i := 0; i < int(inputCountEval); i++ {
|
||||
name, err := getInputName(ortTrainingSession, i, "eval")
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("error retrieving eval input name: %w", err)
|
||||
}
|
||||
evalInputNames[i] = name
|
||||
}
|
||||
|
||||
for i := 0; i < int(outputCountTraining); i++ {
|
||||
name, err := getOutputName(ortTrainingSession, i, "eval")
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("error retrieving eval output name: %w", err)
|
||||
}
|
||||
evalOutputNames[i] = name
|
||||
}
|
||||
}
|
||||
|
||||
return &TrainingInputOutputNames{
|
||||
TrainingInputNames: trainInputNames,
|
||||
EvalInputNames: evalInputNames,
|
||||
TrainingOutputNames: trainOutputNames,
|
||||
EvalOutputNames: evalOutputNames,
|
||||
}, nil
|
||||
}
|
||||
|
||||
// IsTrainingSupported returns true if the training api is supported
|
||||
// by the onnxruntime library.
|
||||
func IsTrainingSupported() bool {
|
||||
return C.IsTrainingApiSupported() != 0
|
||||
}
|
||||
|
||||
func checkTraining() error {
|
||||
if !IsInitialized() {
|
||||
return NotInitializedError
|
||||
}
|
||||
if !IsTrainingSupported() {
|
||||
return trainingNotSupportedError
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func createCCheckpoint(onnxData []byte) (*C.OrtCheckpointState, error) {
|
||||
if e := checkTraining(); e != nil {
|
||||
return nil, e
|
||||
}
|
||||
if len(onnxData) == 0 {
|
||||
return nil, fmt.Errorf("Missing checkpoint data")
|
||||
}
|
||||
var ortCheckpointState *C.OrtCheckpointState
|
||||
status := C.CreateCheckpoint(unsafe.Pointer(&(onnxData[0])), C.size_t(len(onnxData)), &ortCheckpointState)
|
||||
if status != nil {
|
||||
return nil, statusToError(status)
|
||||
}
|
||||
return ortCheckpointState, nil
|
||||
}
|
||||
|
||||
// createCTrainingSessionWithOnnxData creates a C session from byte data using buffers
|
||||
func createCTrainingSessionWithOnnxData(checkpointState *C.OrtCheckpointState,
|
||||
trainingData, evalData, optimizerData []byte,
|
||||
options *SessionOptions) (*C.OrtTrainingSession, error) {
|
||||
if e := checkTraining(); e != nil {
|
||||
return nil, e
|
||||
}
|
||||
|
||||
var ortTrainingSession *C.OrtTrainingSession
|
||||
var ortSessionOptions *C.OrtSessionOptions
|
||||
if options != nil {
|
||||
ortSessionOptions = options.o
|
||||
}
|
||||
|
||||
// eval model is optional
|
||||
var evalDataPtr unsafe.Pointer
|
||||
var evalDataSize C.size_t
|
||||
if len(evalData) > 0 {
|
||||
evalDataPtr = unsafe.Pointer(&(evalData[0]))
|
||||
evalDataSize = C.size_t(len(evalData))
|
||||
}
|
||||
|
||||
// optimizer model is also optional when e.g. getting input and output names
|
||||
var optimizerDataPtr unsafe.Pointer
|
||||
var optimizerDataSize C.size_t
|
||||
if len(optimizerData) > 0 {
|
||||
optimizerDataPtr = unsafe.Pointer(&(optimizerData[0]))
|
||||
optimizerDataSize = C.size_t(len(optimizerData))
|
||||
}
|
||||
|
||||
status := C.CreateTrainingSessionFromBuffer(
|
||||
checkpointState,
|
||||
unsafe.Pointer(&(trainingData[0])), C.size_t(len(trainingData)),
|
||||
evalDataPtr, evalDataSize,
|
||||
optimizerDataPtr, optimizerDataSize,
|
||||
ortEnv, &ortTrainingSession, ortSessionOptions)
|
||||
if status != nil {
|
||||
return nil, statusToError(status)
|
||||
}
|
||||
return ortTrainingSession, nil
|
||||
}
|
||||
|
||||
// createCTrainingSessionWithPaths creates a C session from paths
|
||||
func createCtrainingSessionWithPaths(checkpointState *C.OrtCheckpointState,
|
||||
trainingPath, evalPath, optimizerPath *C.char,
|
||||
options *SessionOptions) (*C.OrtTrainingSession, error) {
|
||||
if e := checkTraining(); e != nil {
|
||||
return nil, e
|
||||
}
|
||||
|
||||
var ortTrainingSession *C.OrtTrainingSession
|
||||
var ortSessionOptions *C.OrtSessionOptions
|
||||
if options != nil {
|
||||
ortSessionOptions = options.o
|
||||
}
|
||||
|
||||
status := C.CreateTrainingSessionFromPaths(checkpointState,
|
||||
trainingPath, evalPath, optimizerPath, ortEnv, &ortTrainingSession, ortSessionOptions)
|
||||
|
||||
if status != nil {
|
||||
return nil, statusToError(status)
|
||||
}
|
||||
return ortTrainingSession, nil
|
||||
}
|
||||
|
||||
// NewTrainingSessionWithOnnxData is like NewTrainingSession, but it accepts
|
||||
// bytes rather than paths to the training assets. Note that there does not
|
||||
// seem to currently be a way to export the trained model from a session
|
||||
// instantiated from bytes. If you wish to export the trained model, you should
|
||||
// use NewTrainingSession instead.
|
||||
func NewTrainingSessionWithOnnxData(checkpointData []byte,
|
||||
trainingData []byte,
|
||||
evalData []byte,
|
||||
optimizerData []byte,
|
||||
inputs, outputs []Value,
|
||||
options *SessionOptions) (*TrainingSession, error) {
|
||||
|
||||
if err := checkTraining(); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
if err := validateInputOutputs(inputs, outputs); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
if len(trainingData) == 0 {
|
||||
return nil, fmt.Errorf("training data has length zero.")
|
||||
}
|
||||
if len(optimizerData) == 0 {
|
||||
return nil, fmt.Errorf("optimizer data has length zero.")
|
||||
}
|
||||
|
||||
// create checkpoint C object
|
||||
ortCheckpointState, e := createCCheckpoint(checkpointData)
|
||||
if e != nil {
|
||||
return nil, fmt.Errorf("error creating C checkpointState: %w", e)
|
||||
}
|
||||
|
||||
// create session C object
|
||||
ortTrainingSession, e := createCTrainingSessionWithOnnxData(ortCheckpointState,
|
||||
trainingData, evalData, optimizerData, options)
|
||||
if e != nil {
|
||||
return nil, fmt.Errorf("error creating C training session: %w", e)
|
||||
}
|
||||
|
||||
inputOrtTensors := make([]*C.OrtValue, len(inputs))
|
||||
outputOrtTensors := make([]*C.OrtValue, len(outputs))
|
||||
for i, v := range inputs {
|
||||
inputOrtTensors[i] = v.GetInternals().ortValue
|
||||
}
|
||||
for i, v := range outputs {
|
||||
outputOrtTensors[i] = v.GetInternals().ortValue
|
||||
}
|
||||
|
||||
return &TrainingSession{
|
||||
ortCheckpointState: ortCheckpointState,
|
||||
ortTrainingSession: ortTrainingSession,
|
||||
inputs: inputOrtTensors,
|
||||
outputs: outputOrtTensors,
|
||||
}, nil
|
||||
}
|
||||
|
||||
func validateInputOutputs(inputs, outputs []Value) error {
|
||||
if len(inputs) == 0 {
|
||||
return fmt.Errorf("inputs must have length greater than zero")
|
||||
}
|
||||
if len(outputs) == 0 {
|
||||
return fmt.Errorf("outputs must have length greater than zero")
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// NewTrainingSession creates a new training session from paths stored on disk.
|
||||
// evalModelPath is optional and can be the empty string. In case it is not
|
||||
// provided, only the checkpoint state can be exported once training is complete
|
||||
// (and not the final inference model).
|
||||
func NewTrainingSession(checkpointStatePath string,
|
||||
trainingModelPath string,
|
||||
evalModelPath string,
|
||||
optimizerModelPath string,
|
||||
inputs, outputs []Value,
|
||||
options *SessionOptions) (*TrainingSession, error) {
|
||||
|
||||
if err := checkTraining(); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
if err := validateInputOutputs(inputs, outputs); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
checkPointContent, e := os.ReadFile(checkpointStatePath)
|
||||
if e != nil {
|
||||
return nil, fmt.Errorf("reading checkpoint data failed: %s", e.Error())
|
||||
}
|
||||
|
||||
// create checkpoint C object
|
||||
ortCheckpointState, e := createCCheckpoint(checkPointContent)
|
||||
if e != nil {
|
||||
return nil, fmt.Errorf("error creating C checkpointState: %w", e)
|
||||
}
|
||||
|
||||
// create session C object
|
||||
if _, err := os.Stat(trainingModelPath); os.IsNotExist(err) {
|
||||
return nil, fmt.Errorf("training model does not exist at path %s", trainingModelPath)
|
||||
}
|
||||
cTrainingPath, err := createOrtCharString(trainingModelPath)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("Error converting training model path to C string: %w", err)
|
||||
}
|
||||
|
||||
if _, err := os.Stat(optimizerModelPath); os.IsNotExist(err) {
|
||||
return nil, fmt.Errorf("optimizer s does not exist at path %s", optimizerModelPath)
|
||||
}
|
||||
cOptimizerPath, err := createOrtCharString(optimizerModelPath)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("Error converting optimizer path to C string: %w", err)
|
||||
}
|
||||
|
||||
// eval is optional
|
||||
var cEvalPath *C.char
|
||||
if evalModelPath != "" {
|
||||
if _, err := os.Stat(evalModelPath); os.IsNotExist(err) {
|
||||
return nil, fmt.Errorf("eval model does not exist at path %s", evalModelPath)
|
||||
}
|
||||
cEvalPath, err = createOrtCharString(evalModelPath)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("Error converting eval path to C string: %w", err)
|
||||
}
|
||||
} else {
|
||||
cEvalPath = nil
|
||||
}
|
||||
|
||||
ortTrainingSession, e := createCtrainingSessionWithPaths(ortCheckpointState,
|
||||
cTrainingPath, cEvalPath, cOptimizerPath, options)
|
||||
if e != nil {
|
||||
return nil, fmt.Errorf("error creating C training session: %w", e)
|
||||
}
|
||||
|
||||
inputOrtTensors := make([]*C.OrtValue, len(inputs))
|
||||
outputOrtTensors := make([]*C.OrtValue, len(outputs))
|
||||
for i, v := range inputs {
|
||||
inputOrtTensors[i] = v.GetInternals().ortValue
|
||||
}
|
||||
for i, v := range outputs {
|
||||
outputOrtTensors[i] = v.GetInternals().ortValue
|
||||
}
|
||||
|
||||
return &TrainingSession{
|
||||
ortCheckpointState: ortCheckpointState,
|
||||
ortTrainingSession: ortTrainingSession,
|
||||
inputs: inputOrtTensors,
|
||||
outputs: outputOrtTensors,
|
||||
evalModelPath: cEvalPath,
|
||||
trainingModelPath: cTrainingPath,
|
||||
optimizerModelPath: cOptimizerPath,
|
||||
}, nil
|
||||
}
|
||||
+547
@@ -0,0 +1,547 @@
|
||||
#include "onnxruntime_wrapper.h"
|
||||
|
||||
static const OrtApi *ort_api = NULL;
|
||||
static const char *ORT_VERSION = NULL;
|
||||
|
||||
static AppendCoreMLProviderFn append_coreml_provider_fn = NULL;
|
||||
|
||||
// The dml_provider_factory.h header for using DirectML is annoying to include
|
||||
// here for a couple reasons:
|
||||
// - It contains C++
|
||||
// - It includes d3d12.h and DirectML.h, both of which may be hard to set up
|
||||
// under mingw
|
||||
// Fortunately, the basic AppendExecutionProvider_DML function from the
|
||||
// OrtDmlApi struct does not rely on any of these things, but we still need the
|
||||
// struct definition itself. Obviously, copying it here is not perfect, and
|
||||
// we'll need to keep an eye on it to make sure it doesn't change between
|
||||
// updates. Most importantly, we need to make sure that the one function we
|
||||
// care about remains at the same place in the struct. Since it's first,
|
||||
// hopefully it's unlikely to change.
|
||||
typedef OrtStatus* (*AppendDirectMLProviderFn)(OrtSessionOptions*, int);
|
||||
typedef struct {
|
||||
AppendDirectMLProviderFn SessionOptionsAppendExecutionProvider_DML;
|
||||
// All of these functions pointers should be irrelevant (and they depend on
|
||||
// other definitions from dml_provider_factory.h), but I'll copy them here
|
||||
// regardless as plain void*s. GetExecutionProviderApi shouldn't write to
|
||||
// this struct anyway, as it only provides a const pointer to it.
|
||||
void *SessionOptionsAppendExecutionProvider_DML1;
|
||||
void *CreateGPUAllocationFromD3DResource;
|
||||
void *FreeGPUAllocation;
|
||||
void *GetD3D12ResourceFromAllocation;
|
||||
void *SessionOptionsAppendExecutionProvider_DML2;
|
||||
} DummyOrtDMLAPI;
|
||||
|
||||
int SetAPIFromBase(OrtApiBase *api_base) {
|
||||
if (!api_base) return 1;
|
||||
ort_api = api_base->GetApi(ORT_API_VERSION);
|
||||
ORT_VERSION = api_base->GetVersionString();
|
||||
if (!ort_api) return 2;
|
||||
return 0;
|
||||
}
|
||||
|
||||
const char *GetVersion() {
|
||||
return ORT_VERSION;
|
||||
}
|
||||
|
||||
void SetCoreMLProviderFunctionPointer(void *ptr) {
|
||||
append_coreml_provider_fn = (AppendCoreMLProviderFn) ptr;
|
||||
}
|
||||
|
||||
void ReleaseOrtStatus(OrtStatus *status) {
|
||||
ort_api->ReleaseStatus(status);
|
||||
}
|
||||
|
||||
OrtStatus *CreateOrtEnv(char *name, OrtEnv **env) {
|
||||
return ort_api->CreateEnv(ORT_LOGGING_LEVEL_ERROR, name, env);
|
||||
}
|
||||
|
||||
OrtStatus *DisableTelemetry(OrtEnv *env) {
|
||||
return ort_api->DisableTelemetryEvents(env);
|
||||
}
|
||||
|
||||
OrtStatus *EnableTelemetry(OrtEnv *env) {
|
||||
return ort_api->EnableTelemetryEvents(env);
|
||||
}
|
||||
|
||||
void ReleaseOrtEnv(OrtEnv *env) {
|
||||
ort_api->ReleaseEnv(env);
|
||||
}
|
||||
|
||||
OrtStatus *CreateOrtMemoryInfo(OrtMemoryInfo **mem_info) {
|
||||
return ort_api->CreateCpuMemoryInfo(OrtArenaAllocator, OrtMemTypeDefault,
|
||||
mem_info);
|
||||
}
|
||||
|
||||
void ReleaseOrtMemoryInfo(OrtMemoryInfo *info) {
|
||||
ort_api->ReleaseMemoryInfo(info);
|
||||
}
|
||||
|
||||
const char *GetErrorMessage(OrtStatus *status) {
|
||||
if (!status) return "No error (NULL status)";
|
||||
return ort_api->GetErrorMessage(status);
|
||||
}
|
||||
|
||||
OrtStatus *CreateSessionOptions(OrtSessionOptions **o) {
|
||||
return ort_api->CreateSessionOptions(o);
|
||||
}
|
||||
|
||||
void ReleaseSessionOptions(OrtSessionOptions *o) {
|
||||
ort_api->ReleaseSessionOptions(o);
|
||||
}
|
||||
|
||||
OrtStatus *SetIntraOpNumThreads(OrtSessionOptions *o, int n) {
|
||||
return ort_api->SetIntraOpNumThreads(o, n);
|
||||
}
|
||||
|
||||
OrtStatus *SetInterOpNumThreads(OrtSessionOptions *o, int n) {
|
||||
return ort_api->SetInterOpNumThreads(o, n);
|
||||
}
|
||||
|
||||
OrtStatus *SetCpuMemArena(OrtSessionOptions *o, int use_arena){
|
||||
if (use_arena)
|
||||
return ort_api->EnableCpuMemArena(o);
|
||||
return ort_api->DisableCpuMemArena(o);
|
||||
}
|
||||
|
||||
OrtStatus *SetMemPattern(OrtSessionOptions *o, int use_mem_pattern){
|
||||
if (use_mem_pattern)
|
||||
return ort_api->EnableMemPattern(o);
|
||||
return ort_api->DisableMemPattern(o);
|
||||
}
|
||||
|
||||
OrtStatus *AppendExecutionProviderCUDAV2(OrtSessionOptions *o,
|
||||
OrtCUDAProviderOptionsV2 *cuda_options) {
|
||||
return ort_api->SessionOptionsAppendExecutionProvider_CUDA_V2(o,
|
||||
cuda_options);
|
||||
}
|
||||
|
||||
OrtStatus *CreateCUDAProviderOptions(OrtCUDAProviderOptionsV2 **o) {
|
||||
return ort_api->CreateCUDAProviderOptions(o);
|
||||
}
|
||||
|
||||
void ReleaseCUDAProviderOptions(OrtCUDAProviderOptionsV2 *o) {
|
||||
ort_api->ReleaseCUDAProviderOptions(o);
|
||||
}
|
||||
|
||||
OrtStatus *UpdateCUDAProviderOptions(OrtCUDAProviderOptionsV2 *o,
|
||||
const char **keys, const char **values, int num_keys) {
|
||||
return ort_api->UpdateCUDAProviderOptions(o, keys, values, num_keys);
|
||||
}
|
||||
|
||||
OrtStatus *CreateTensorRTProviderOptions(OrtTensorRTProviderOptionsV2 **o) {
|
||||
return ort_api->CreateTensorRTProviderOptions(o);
|
||||
}
|
||||
|
||||
void ReleaseTensorRTProviderOptions(OrtTensorRTProviderOptionsV2 *o) {
|
||||
ort_api->ReleaseTensorRTProviderOptions(o);
|
||||
}
|
||||
|
||||
OrtStatus *UpdateTensorRTProviderOptions(OrtTensorRTProviderOptionsV2 *o,
|
||||
const char **keys, const char **values, int num_keys) {
|
||||
return ort_api->UpdateTensorRTProviderOptions(o, keys, values, num_keys);
|
||||
}
|
||||
|
||||
OrtStatus *AppendExecutionProviderTensorRTV2(OrtSessionOptions *o,
|
||||
OrtTensorRTProviderOptionsV2 *tensor_rt_options) {
|
||||
return ort_api->SessionOptionsAppendExecutionProvider_TensorRT_V2(o,
|
||||
tensor_rt_options);
|
||||
}
|
||||
|
||||
OrtStatus *AppendExecutionProviderCoreML(OrtSessionOptions *o,
|
||||
uint32_t flags) {
|
||||
if (!append_coreml_provider_fn) {
|
||||
return ort_api->CreateStatus(ORT_NOT_IMPLEMENTED, "Your platform or "
|
||||
"onnxruntime library does not support CoreML");
|
||||
}
|
||||
return append_coreml_provider_fn(o, flags);
|
||||
}
|
||||
|
||||
OrtStatus *AppendExecutionProviderDirectML(OrtSessionOptions *o,
|
||||
int device_id) {
|
||||
DummyOrtDMLAPI *dml_api = NULL;
|
||||
OrtStatus *status = NULL;
|
||||
status = ort_api->GetExecutionProviderApi("DML", ORT_API_VERSION,
|
||||
(const void **) (&dml_api));
|
||||
if (status) return status;
|
||||
status = dml_api->SessionOptionsAppendExecutionProvider_DML(o, device_id);
|
||||
return status;
|
||||
}
|
||||
|
||||
OrtStatus *AppendExecutionProviderOpenVINOV2(OrtSessionOptions *o,
|
||||
const char **keys, const char **values, int num_keys) {
|
||||
return ort_api->SessionOptionsAppendExecutionProvider_OpenVINO_V2(o, keys,
|
||||
values, num_keys);
|
||||
}
|
||||
|
||||
OrtStatus *CreateSession(void *model_data, size_t model_data_length,
|
||||
OrtEnv *env, OrtSession **out, OrtSessionOptions *options) {
|
||||
OrtStatus *status = NULL;
|
||||
int default_options = 0;
|
||||
if (!options) {
|
||||
default_options = 1;
|
||||
status = ort_api->CreateSessionOptions(&options);
|
||||
if (status) return status;
|
||||
}
|
||||
status = ort_api->CreateSessionFromArray(env, model_data, model_data_length,
|
||||
options, out);
|
||||
if (default_options) {
|
||||
// If we created a default, empty, options struct, we don't need to keep it
|
||||
// after creating the session.
|
||||
ort_api->ReleaseSessionOptions(options);
|
||||
}
|
||||
return status;
|
||||
}
|
||||
|
||||
OrtStatus *CreateSessionFromFile(char *model_path, OrtEnv *env,
|
||||
OrtSession **out, OrtSessionOptions *options) {
|
||||
// Nearly identical to CreateSession, except invokes ort_api->CreateSession
|
||||
// rather than ort_api->CreateSessionFromArray.
|
||||
OrtStatus *status = NULL;
|
||||
int default_options = 0;
|
||||
if (!options) {
|
||||
default_options = 1;
|
||||
status = ort_api->CreateSessionOptions(&options);
|
||||
if (status) return status;
|
||||
}
|
||||
status = ort_api->CreateSession(env, (const ORTCHAR_T*) model_path, options,
|
||||
out);
|
||||
if (default_options) ort_api->ReleaseSessionOptions(options);
|
||||
return status;
|
||||
}
|
||||
|
||||
OrtStatus *RunOrtSession(OrtSession *session,
|
||||
OrtValue **inputs, char **input_names, int input_count,
|
||||
OrtValue **outputs, char **output_names, int output_count) {
|
||||
OrtStatus *status = NULL;
|
||||
status = ort_api->Run(session, NULL, (const char* const*) input_names,
|
||||
(const OrtValue* const*) inputs, input_count,
|
||||
(const char* const*) output_names, output_count, outputs);
|
||||
return status;
|
||||
}
|
||||
|
||||
void ReleaseOrtSession(OrtSession *session) {
|
||||
ort_api->ReleaseSession(session);
|
||||
}
|
||||
|
||||
OrtStatus *SessionGetInputCount(OrtSession *session, size_t *result) {
|
||||
return ort_api->SessionGetInputCount(session, result);
|
||||
}
|
||||
|
||||
OrtStatus *SessionGetOutputCount(OrtSession *session, size_t *result) {
|
||||
return ort_api->SessionGetOutputCount(session, result);
|
||||
}
|
||||
|
||||
void ReleaseOrtValue(OrtValue *value) {
|
||||
ort_api->ReleaseValue(value);
|
||||
}
|
||||
|
||||
OrtStatus *CreateOrtTensorWithShape(void *data, size_t data_size,
|
||||
int64_t *shape, int64_t shape_size, OrtMemoryInfo *mem_info,
|
||||
ONNXTensorElementDataType dtype, OrtValue **out) {
|
||||
OrtStatus *status = NULL;
|
||||
status = ort_api->CreateTensorWithDataAsOrtValue(mem_info, data, data_size,
|
||||
shape, shape_size, dtype, out);
|
||||
return status;
|
||||
}
|
||||
|
||||
OrtStatus *GetTensorTypeAndShape(const OrtValue *value, OrtTensorTypeAndShapeInfo **out) {
|
||||
return ort_api->GetTensorTypeAndShape(value, out);
|
||||
}
|
||||
|
||||
OrtStatus *GetDimensionsCount(const OrtTensorTypeAndShapeInfo *info, size_t *out) {
|
||||
return ort_api->GetDimensionsCount(info, out);
|
||||
}
|
||||
|
||||
OrtStatus *GetDimensions(const OrtTensorTypeAndShapeInfo *info, int64_t *dim_values, size_t dim_values_length) {
|
||||
return ort_api->GetDimensions(info, dim_values, dim_values_length);
|
||||
}
|
||||
|
||||
OrtStatus *GetTensorElementType(const OrtTensorTypeAndShapeInfo *info, enum ONNXTensorElementDataType *out) {
|
||||
return ort_api->GetTensorElementType(info, out);
|
||||
}
|
||||
|
||||
void ReleaseTensorTypeAndShapeInfo(OrtTensorTypeAndShapeInfo *input) {
|
||||
ort_api->ReleaseTensorTypeAndShapeInfo(input);
|
||||
}
|
||||
|
||||
OrtStatus *GetTensorMutableData(OrtValue *value, void **out) {
|
||||
return ort_api->GetTensorMutableData(value, out);
|
||||
}
|
||||
|
||||
OrtStatus *SessionGetInputName(OrtSession *session, size_t i, char **name) {
|
||||
OrtAllocator *allocator = NULL;
|
||||
OrtStatus *status = NULL;
|
||||
status = ort_api->GetAllocatorWithDefaultOptions(&allocator);
|
||||
if (status) return status;
|
||||
return ort_api->SessionGetInputName(session, i, allocator, name);
|
||||
}
|
||||
|
||||
OrtStatus *SessionGetOutputName(OrtSession *session, size_t i, char **name) {
|
||||
OrtAllocator *allocator = NULL;
|
||||
OrtStatus *status = NULL;
|
||||
status = ort_api->GetAllocatorWithDefaultOptions(&allocator);
|
||||
if (status) return status;
|
||||
return ort_api->SessionGetOutputName(session, i, allocator, name);
|
||||
}
|
||||
|
||||
OrtStatus *FreeWithDefaultORTAllocator(void *to_free) {
|
||||
OrtAllocator *allocator = NULL;
|
||||
OrtStatus *status = NULL;
|
||||
status = ort_api->GetAllocatorWithDefaultOptions(&allocator);
|
||||
if (status) return status;
|
||||
return ort_api->AllocatorFree(allocator, to_free);
|
||||
}
|
||||
|
||||
OrtStatus *SessionGetInputTypeInfo(OrtSession *session, size_t i,
|
||||
OrtTypeInfo **out) {
|
||||
return ort_api->SessionGetInputTypeInfo(session, i, out);
|
||||
}
|
||||
|
||||
OrtStatus *SessionGetOutputTypeInfo(OrtSession *session, size_t i,
|
||||
OrtTypeInfo **out) {
|
||||
return ort_api->SessionGetOutputTypeInfo(session, i, out);
|
||||
}
|
||||
|
||||
void ReleaseTypeInfo(OrtTypeInfo *o) {
|
||||
ort_api->ReleaseTypeInfo(o);
|
||||
}
|
||||
|
||||
OrtStatus *GetONNXTypeFromTypeInfo(OrtTypeInfo *info, enum ONNXType *out) {
|
||||
return ort_api->GetOnnxTypeFromTypeInfo(info, out);
|
||||
}
|
||||
|
||||
OrtStatus *CastTypeInfoToTensorInfo(OrtTypeInfo *type_info,
|
||||
OrtTensorTypeAndShapeInfo **out) {
|
||||
return ort_api->CastTypeInfoToTensorInfo(type_info,
|
||||
(const OrtTensorTypeAndShapeInfo **) out);
|
||||
}
|
||||
|
||||
OrtStatus *SessionGetModelMetadata(OrtSession *s, OrtModelMetadata **m) {
|
||||
return ort_api->SessionGetModelMetadata(s, m);
|
||||
}
|
||||
|
||||
void ReleaseModelMetadata(OrtModelMetadata *m) {
|
||||
return ort_api->ReleaseModelMetadata(m);
|
||||
}
|
||||
|
||||
OrtStatus *ModelMetadataGetProducerName(OrtModelMetadata *m, char **name) {
|
||||
OrtAllocator *allocator = NULL;
|
||||
OrtStatus *status = NULL;
|
||||
status = ort_api->GetAllocatorWithDefaultOptions(&allocator);
|
||||
if (status) return status;
|
||||
return ort_api->ModelMetadataGetProducerName(m, allocator, name);
|
||||
}
|
||||
|
||||
OrtStatus *ModelMetadataGetGraphName(OrtModelMetadata *m, char **name) {
|
||||
OrtAllocator *allocator = NULL;
|
||||
OrtStatus *status = NULL;
|
||||
status = ort_api->GetAllocatorWithDefaultOptions(&allocator);
|
||||
if (status) return status;
|
||||
return ort_api->ModelMetadataGetGraphName(m, allocator, name);
|
||||
}
|
||||
|
||||
OrtStatus *ModelMetadataGetDomain(OrtModelMetadata *m, char **domain) {
|
||||
OrtAllocator *allocator = NULL;
|
||||
OrtStatus *status = NULL;
|
||||
status = ort_api->GetAllocatorWithDefaultOptions(&allocator);
|
||||
if (status) return status;
|
||||
return ort_api->ModelMetadataGetDomain(m, allocator, domain);
|
||||
}
|
||||
|
||||
OrtStatus *ModelMetadataGetDescription(OrtModelMetadata *m, char **desc) {
|
||||
OrtAllocator *allocator = NULL;
|
||||
OrtStatus *status = NULL;
|
||||
status = ort_api->GetAllocatorWithDefaultOptions(&allocator);
|
||||
if (status) return status;
|
||||
return ort_api->ModelMetadataGetDescription(m, allocator, desc);
|
||||
}
|
||||
|
||||
OrtStatus *ModelMetadataLookupCustomMetadataMap(OrtModelMetadata *m, char *key,
|
||||
char **value) {
|
||||
OrtAllocator *allocator = NULL;
|
||||
OrtStatus *status = NULL;
|
||||
status = ort_api->GetAllocatorWithDefaultOptions(&allocator);
|
||||
if (status) return status;
|
||||
return ort_api->ModelMetadataLookupCustomMetadataMap(m, allocator, key,
|
||||
value);
|
||||
}
|
||||
|
||||
OrtStatus *ModelMetadataGetCustomMetadataMapKeys(OrtModelMetadata *m,
|
||||
char ***keys, int64_t *num_keys) {
|
||||
OrtAllocator *allocator = NULL;
|
||||
OrtStatus *status = NULL;
|
||||
status = ort_api->GetAllocatorWithDefaultOptions(&allocator);
|
||||
if (status) return status;
|
||||
return ort_api->ModelMetadataGetCustomMetadataMapKeys(m, allocator, keys,
|
||||
num_keys);
|
||||
}
|
||||
|
||||
OrtStatus *ModelMetadataGetVersion(OrtModelMetadata *m, int64_t *version) {
|
||||
return ort_api->ModelMetadataGetVersion(m, version);
|
||||
}
|
||||
|
||||
OrtStatus *GetValue(OrtValue *container, int index, OrtValue **dst) {
|
||||
OrtAllocator *allocator = NULL;
|
||||
OrtStatus *status = NULL;
|
||||
status = ort_api->GetAllocatorWithDefaultOptions(&allocator);
|
||||
if (status) return status;
|
||||
return ort_api->GetValue(container, index, allocator, dst);
|
||||
}
|
||||
|
||||
OrtStatus *GetValueType(OrtValue *v, enum ONNXType *out) {
|
||||
return ort_api->GetValueType(v, out);
|
||||
}
|
||||
|
||||
OrtStatus *GetValueCount(OrtValue *v, size_t *out) {
|
||||
return ort_api->GetValueCount(v, out);
|
||||
}
|
||||
|
||||
OrtStatus *CreateOrtValue(OrtValue **in, size_t num_values,
|
||||
enum ONNXType value_type, OrtValue **out) {
|
||||
return ort_api->CreateValue((const OrtValue* const*) in, num_values,
|
||||
value_type, out);
|
||||
}
|
||||
|
||||
// TRAINING API WRAPPER
|
||||
|
||||
static const OrtTrainingApi *ort_training_api = NULL;
|
||||
|
||||
void SetTrainingApi() {
|
||||
ort_training_api = ort_api->GetTrainingApi(ORT_API_VERSION);
|
||||
}
|
||||
|
||||
int IsTrainingApiSupported() {
|
||||
return ort_training_api != NULL;
|
||||
}
|
||||
|
||||
OrtStatus *CreateCheckpoint(void *checkpoint_data, size_t checkpoint_data_length, OrtCheckpointState **out) {
|
||||
OrtStatus *status = NULL;
|
||||
status = ort_training_api->LoadCheckpointFromBuffer(checkpoint_data, checkpoint_data_length, out);
|
||||
return status;
|
||||
}
|
||||
|
||||
OrtStatus *CreateTrainingSessionFromBuffer(OrtCheckpointState *checkpoint_state,
|
||||
void *training_model_data, size_t training_model_data_length,
|
||||
void *eval_model_data, size_t eval_model_data_length,
|
||||
void *optim_model_data, size_t optim_model_data_length,
|
||||
OrtEnv *env, OrtTrainingSession **out, OrtSessionOptions *options) {
|
||||
OrtStatus *status = NULL;
|
||||
int default_options = 0;
|
||||
if (!options) {
|
||||
default_options = 1;
|
||||
status = ort_api->CreateSessionOptions(&options);
|
||||
if (status) return status;
|
||||
}
|
||||
status = ort_training_api->CreateTrainingSessionFromBuffer(env, options, checkpoint_state,
|
||||
training_model_data, training_model_data_length, eval_model_data, eval_model_data_length,
|
||||
optim_model_data, optim_model_data_length, out);
|
||||
if (default_options) {
|
||||
ort_api->ReleaseSessionOptions(options);
|
||||
}
|
||||
return status;
|
||||
}
|
||||
|
||||
OrtStatus *CreateTrainingSessionFromPaths(OrtCheckpointState *checkpoint_state,
|
||||
char *training_model_path, char *eval_model_path, char *optim_model_path,
|
||||
OrtEnv *env, OrtTrainingSession **out, OrtSessionOptions *options) {
|
||||
OrtStatus *status = NULL;
|
||||
int default_options = 0;
|
||||
if (!options) {
|
||||
default_options = 1;
|
||||
status = ort_api->CreateSessionOptions(&options);
|
||||
if (status) return status;
|
||||
}
|
||||
status = ort_training_api->CreateTrainingSession(env, options, checkpoint_state,
|
||||
training_model_path, eval_model_path, optim_model_path, out);
|
||||
if (default_options) {
|
||||
ort_api->ReleaseSessionOptions(options);
|
||||
}
|
||||
return status;
|
||||
}
|
||||
|
||||
OrtStatus *TrainingSessionGetInputCount(OrtTrainingSession *training_session, size_t *result_training, size_t *result_eval) {
|
||||
OrtStatus *status = NULL;
|
||||
status = ort_training_api->TrainingSessionGetTrainingModelInputCount(training_session, result_training);
|
||||
if (status) return status;
|
||||
status = ort_training_api->TrainingSessionGetEvalModelInputCount(training_session, result_eval);
|
||||
return status;
|
||||
}
|
||||
|
||||
OrtStatus *TrainingSessionGetOutputCount(OrtTrainingSession *training_session, size_t *result_training, size_t *result_eval) {
|
||||
OrtStatus *status = NULL;
|
||||
status = ort_training_api->TrainingSessionGetTrainingModelOutputCount(training_session, result_training);
|
||||
if (status) return status;
|
||||
status = ort_training_api->TrainingSessionGetEvalModelOutputCount(training_session, result_eval);
|
||||
return status;
|
||||
}
|
||||
|
||||
OrtStatus *TrainingSessionGetTrainingInputName(OrtTrainingSession *training_session, size_t i, char **name) {
|
||||
OrtAllocator *allocator = NULL;
|
||||
OrtStatus *status = NULL;
|
||||
status = ort_api->GetAllocatorWithDefaultOptions(&allocator);
|
||||
if (status) return status;
|
||||
return ort_training_api->TrainingSessionGetTrainingModelInputName(training_session, i, allocator, name);
|
||||
}
|
||||
|
||||
OrtStatus *TrainingSessionGetTrainingOutputName(OrtTrainingSession *training_session, size_t i, char **name) {
|
||||
OrtAllocator *allocator = NULL;
|
||||
OrtStatus *status = NULL;
|
||||
status = ort_api->GetAllocatorWithDefaultOptions(&allocator);
|
||||
if (status) return status;
|
||||
return ort_training_api->TrainingSessionGetTrainingModelOutputName(training_session, i, allocator, name);
|
||||
}
|
||||
|
||||
OrtStatus *TrainingSessionGetEvalInputName(OrtTrainingSession *training_session, size_t i, char **name) {
|
||||
OrtAllocator *allocator = NULL;
|
||||
OrtStatus *status = NULL;
|
||||
status = ort_api->GetAllocatorWithDefaultOptions(&allocator);
|
||||
if (status) return status;
|
||||
return ort_training_api->TrainingSessionGetEvalModelInputName(training_session, i, allocator, name);
|
||||
}
|
||||
|
||||
OrtStatus *TrainingSessionGetEvalOutputName(OrtTrainingSession *training_session, size_t i, char **name) {
|
||||
OrtAllocator *allocator = NULL;
|
||||
OrtStatus *status = NULL;
|
||||
status = ort_api->GetAllocatorWithDefaultOptions(&allocator);
|
||||
if (status) return status;
|
||||
return ort_training_api->TrainingSessionGetEvalModelOutputName(training_session, i, allocator, name);
|
||||
}
|
||||
|
||||
OrtStatus *TrainStep(OrtTrainingSession *training_session, size_t inputs_len, OrtValue **inputs, size_t output_len, OrtValue **outputs) {
|
||||
OrtStatus *status = NULL;
|
||||
status = ort_training_api->TrainStep(training_session, NULL, inputs_len, (const OrtValue* const*) inputs, output_len, outputs);
|
||||
return status;
|
||||
}
|
||||
|
||||
OrtStatus *OptimizerStep(OrtTrainingSession *training_session) {
|
||||
OrtStatus *status = NULL;
|
||||
status = ort_training_api->OptimizerStep(training_session, NULL);
|
||||
return status;
|
||||
}
|
||||
|
||||
OrtStatus *LazyResetGrad(OrtTrainingSession *training_session) {
|
||||
OrtStatus *status = NULL;
|
||||
status = ort_training_api->LazyResetGrad(training_session);
|
||||
return status;
|
||||
}
|
||||
|
||||
OrtStatus *SaveCheckpoint(OrtCheckpointState *checkpoint, char *path, size_t include_optimizer) {
|
||||
OrtStatus *status = NULL;
|
||||
status = ort_training_api->SaveCheckpoint(checkpoint, path, include_optimizer);
|
||||
return status;
|
||||
}
|
||||
|
||||
OrtStatus *ExportModel(OrtTrainingSession *training_session, char *path, size_t outputs_len, char **output_names) {
|
||||
OrtStatus *status = NULL;
|
||||
status = ort_training_api->ExportModelForInferencing(training_session, path, outputs_len, (const char* const*) output_names);
|
||||
return status;
|
||||
}
|
||||
|
||||
void ReleaseOrtTrainingSession(OrtTrainingSession *session) {
|
||||
ort_training_api->ReleaseTrainingSession(session);
|
||||
}
|
||||
|
||||
void ReleaseCheckpointState(OrtCheckpointState *checkpoint) {
|
||||
ort_training_api->ReleaseCheckpointState(checkpoint);
|
||||
}
|
||||
|
||||
+329
@@ -0,0 +1,329 @@
|
||||
#ifndef ONNXRUNTIME_WRAPPER_H
|
||||
#define ONNXRUNTIME_WRAPPER_H
|
||||
|
||||
// We want to always use the unix-like onnxruntime C APIs, even on Windows, so
|
||||
// we need to undefine _WIN32 before including onnxruntime_c_api.h. However,
|
||||
// this requires a careful song-and-dance.
|
||||
|
||||
// First, include these common headers, as they get transitively included by
|
||||
// onnxruntime_c_api.h. We need to include them ourselves, first, so that the
|
||||
// preprocessor will skip them while _WIN32 is undefined.
|
||||
#include <stdio.h>
|
||||
#include <stdlib.h>
|
||||
#include <string.h>
|
||||
|
||||
// Next, we actually include the header.
|
||||
#undef _WIN32
|
||||
#include "onnxruntime_c_api.h"
|
||||
#include "onnxruntime_training_c_api.h"
|
||||
|
||||
// ... However, mingw will complain if _WIN32 is *not* defined! So redefine it.
|
||||
#define _WIN32
|
||||
|
||||
#ifdef __cplusplus
|
||||
extern "C" {
|
||||
#endif
|
||||
|
||||
// Used for the OrtSessionOptionsAppendExecutionProvider_CoreML function
|
||||
// pointer on supported systems. Must match the signature in
|
||||
// coreml_provider_factory.h provided along with the onnxruntime releases for
|
||||
// Apple platforms.
|
||||
typedef OrtStatus* (*AppendCoreMLProviderFn)(OrtSessionOptions*, uint32_t);
|
||||
|
||||
// Takes a pointer to the api_base struct in order to obtain the OrtApi
|
||||
// pointer. Intended to be called from Go. Returns nonzero on error.
|
||||
int SetAPIFromBase(OrtApiBase *api_base);
|
||||
|
||||
// Get the version of the Onnxruntime library for logging.
|
||||
const char *GetVersion();
|
||||
|
||||
// OrtSessionOptionsAppendExecutionProvider_CoreML is exported directly from
|
||||
// the Apple .dylib, so we call this function on Apple platforms to set the
|
||||
// function pointer to the correct address. On other platforms, the function
|
||||
// pointer should remain NULL.
|
||||
void SetCoreMLProviderFunctionPointer(void *ptr);
|
||||
|
||||
// Wraps ort_api->ReleaseStatus(status)
|
||||
void ReleaseOrtStatus(OrtStatus *status);
|
||||
|
||||
// Wraps calling ort_api->CreateEnv. Returns a non-NULL status on error.
|
||||
OrtStatus *CreateOrtEnv(char *name, OrtEnv **env);
|
||||
|
||||
// Wraps ort_api->DisableTelemetryEvents. Returns a non-NULL status on error.
|
||||
OrtStatus *DisableTelemetry(OrtEnv *env);
|
||||
|
||||
// Wraps ort_api->EnableTelemetryEvents. Returns a non-NULL status on error.
|
||||
OrtStatus *EnableTelemetry(OrtEnv *env);
|
||||
|
||||
// Wraps ort_api->ReleaseEnv
|
||||
void ReleaseOrtEnv(OrtEnv *env);
|
||||
|
||||
// Wraps ort_api->CreateCpuMemoryInfo with some basic, default settings.
|
||||
OrtStatus *CreateOrtMemoryInfo(OrtMemoryInfo **mem_info);
|
||||
|
||||
// Wraps ort_api->ReleaseMemoryInfo
|
||||
void ReleaseOrtMemoryInfo(OrtMemoryInfo *info);
|
||||
|
||||
// Returns the message associated with the given ORT status.
|
||||
const char *GetErrorMessage(OrtStatus *status);
|
||||
|
||||
// Wraps ort_api->CreateSessionOptions
|
||||
OrtStatus *CreateSessionOptions(OrtSessionOptions **o);
|
||||
|
||||
// Wraps ort_api->ReleaseSessionOptions
|
||||
void ReleaseSessionOptions(OrtSessionOptions *o);
|
||||
|
||||
// Wraps ort_api->SetIntraOpNumThreads
|
||||
OrtStatus *SetIntraOpNumThreads(OrtSessionOptions *o, int n);
|
||||
|
||||
// Wraps ort_api->SetInterOpNumThreads
|
||||
OrtStatus *SetInterOpNumThreads(OrtSessionOptions *o, int n);
|
||||
|
||||
// Wraps ort_api->EnableCpuMemArena & ort_api->DisableCpuMemArena
|
||||
OrtStatus *SetCpuMemArena(OrtSessionOptions *o, int use_arena);
|
||||
|
||||
// Wraps ort_api->EnableMemPattern & ort_api->DisableMemPattern
|
||||
OrtStatus *SetMemPattern(OrtSessionOptions *o, int use_mem_pattern);
|
||||
|
||||
// Wraps ort_api->SessionOptionsAppendExecutionProvider_CUDA_V2
|
||||
OrtStatus *AppendExecutionProviderCUDAV2(OrtSessionOptions *o,
|
||||
OrtCUDAProviderOptionsV2 *cuda_options);
|
||||
|
||||
// Wraps ort_api->CreateCUDAProviderOptions
|
||||
OrtStatus *CreateCUDAProviderOptions(OrtCUDAProviderOptionsV2 **o);
|
||||
|
||||
// Wraps ort_api->ReleaseCUDAProviderOptions
|
||||
void ReleaseCUDAProviderOptions(OrtCUDAProviderOptionsV2 *o);
|
||||
|
||||
// Wraps ort_api->UpdateCUDAProviderOptions
|
||||
OrtStatus *UpdateCUDAProviderOptions(OrtCUDAProviderOptionsV2 *o,
|
||||
const char **keys, const char **values, int num_keys);
|
||||
|
||||
// Wraps ort_api->CreateTensorRTProviderOptions
|
||||
OrtStatus *CreateTensorRTProviderOptions(OrtTensorRTProviderOptionsV2 **o);
|
||||
|
||||
// Wraps ort_api->ReleaseTensorRTProviderOptions
|
||||
void ReleaseTensorRTProviderOptions(OrtTensorRTProviderOptionsV2 *o);
|
||||
|
||||
// Wraps ort_api->UpdateTensorRTProviderOptions
|
||||
OrtStatus *UpdateTensorRTProviderOptions(OrtTensorRTProviderOptionsV2 *o,
|
||||
const char **keys, const char **values, int num_keys);
|
||||
|
||||
// Wraps ort_api->SessionOptionsAppendExecutionProvider_TensorRT_V2
|
||||
OrtStatus *AppendExecutionProviderTensorRTV2(OrtSessionOptions *o,
|
||||
OrtTensorRTProviderOptionsV2 *tensor_rt_options);
|
||||
|
||||
// Wraps OrtSessionOptionsAppendExecutionProvider_CoreML, exported from the
|
||||
// dylib on Apple devices. Safely returns a non-NULL status on other platforms.
|
||||
OrtStatus *AppendExecutionProviderCoreML(OrtSessionOptions *o,
|
||||
uint32_t flags);
|
||||
|
||||
// Wraps getting the OrtDmlApi struct and calling
|
||||
// dml_api->SessionOptionsAppendExecutionProvider_DML.
|
||||
OrtStatus *AppendExecutionProviderDirectML(OrtSessionOptions *o,
|
||||
int device_id);
|
||||
|
||||
// Wraps ort_api->AppendExecutionProvider_OpenVINO_V2
|
||||
OrtStatus *AppendExecutionProviderOpenVINOV2(OrtSessionOptions *o,
|
||||
const char **keys, const char **values, int num_keys);
|
||||
|
||||
// Creates an ORT session using the given model. The given options pointer may
|
||||
// be NULL; if it is, then we'll use default options.
|
||||
OrtStatus *CreateSession(void *model_data, size_t model_data_length,
|
||||
OrtEnv *env, OrtSession **out, OrtSessionOptions *options);
|
||||
|
||||
// Like the CreateSession function, but takes a path to a model rather than a
|
||||
// buffer containing it.
|
||||
OrtStatus *CreateSessionFromFile(char *model_path, OrtEnv *env,
|
||||
OrtSession **out, OrtSessionOptions *options);
|
||||
|
||||
// Runs an ORT session with the given input and output tensors, along with
|
||||
// their names. In our use case, outputs must NOT be NULL.
|
||||
OrtStatus *RunOrtSession(OrtSession *session,
|
||||
OrtValue **inputs, char **input_names, int input_count,
|
||||
OrtValue **outputs, char **output_names, int output_count);
|
||||
|
||||
// Wraps ort_api->ReleaseSession
|
||||
void ReleaseOrtSession(OrtSession *session);
|
||||
|
||||
// Wraps ort_api->SessionGetInputCount.
|
||||
OrtStatus *SessionGetInputCount(OrtSession *session, size_t *result);
|
||||
|
||||
// Wraps ort_api->SessionGetOutputCount.
|
||||
OrtStatus *SessionGetOutputCount(OrtSession *session, size_t *result);
|
||||
|
||||
// Used to free OrtValue instances, such as tensors.
|
||||
void ReleaseOrtValue(OrtValue *value);
|
||||
|
||||
// Creates an OrtValue tensor with the given shape, and backed by the user-
|
||||
// supplied data buffer.
|
||||
OrtStatus *CreateOrtTensorWithShape(void *data, size_t data_size,
|
||||
int64_t *shape, int64_t shape_size, OrtMemoryInfo *mem_info,
|
||||
ONNXTensorElementDataType dtype, OrtValue **out);
|
||||
|
||||
// Wraps ort_api->GetTensorTypeAndShape
|
||||
OrtStatus *GetTensorTypeAndShape(const OrtValue *value, OrtTensorTypeAndShapeInfo **out);
|
||||
|
||||
// Wraps ort_api->GetDimensionsCount
|
||||
OrtStatus *GetDimensionsCount(const OrtTensorTypeAndShapeInfo *info, size_t *out);
|
||||
|
||||
// Wraps ort_api->GetDimensions
|
||||
OrtStatus *GetDimensions(const OrtTensorTypeAndShapeInfo *info, int64_t *dim_values, size_t dim_values_length);
|
||||
|
||||
// Wraps ort_api->GetTensorElementType
|
||||
OrtStatus *GetTensorElementType(const OrtTensorTypeAndShapeInfo *info, enum ONNXTensorElementDataType *out);
|
||||
|
||||
// Wraps ort_api->ReleaseTensorTypeAndShapeInfo
|
||||
void ReleaseTensorTypeAndShapeInfo(OrtTensorTypeAndShapeInfo *input);
|
||||
|
||||
// Wraps ort_api->GetTensorMutableData
|
||||
OrtStatus *GetTensorMutableData(OrtValue *value, void **out);
|
||||
|
||||
// Wraps ort_api->SessionGetInputName, using the default allocator.
|
||||
OrtStatus *SessionGetInputName(OrtSession *session, size_t i, char **name);
|
||||
|
||||
// Wraps ort_api->SessionGetOutputName, using the default allocator.
|
||||
OrtStatus *SessionGetOutputName(OrtSession *session, size_t i, char **name);
|
||||
|
||||
// Frees anything that was allocated using the default ORT allocator.
|
||||
OrtStatus *FreeWithDefaultORTAllocator(void *to_free);
|
||||
|
||||
// Wraps ort_api->SessionGetInputTypeInfo.
|
||||
OrtStatus *SessionGetInputTypeInfo(OrtSession *session, size_t i,
|
||||
OrtTypeInfo **out);
|
||||
|
||||
// Wraps ort_api->SessionGetOutputTypeInfo.
|
||||
OrtStatus *SessionGetOutputTypeInfo(OrtSession *session, size_t i,
|
||||
OrtTypeInfo **out);
|
||||
|
||||
// If the type_info is for a tensor, sets out to the a pointer to the tensor's
|
||||
// NameAndTypeInfo. Do _not_ free the out pointer; it will be freed when
|
||||
// type_info is released.
|
||||
//
|
||||
// Wraps ort_api->CastTypeInfoToTensorInfo.
|
||||
OrtStatus *CastTypeInfoToTensorInfo(OrtTypeInfo *type_info,
|
||||
OrtTensorTypeAndShapeInfo **out);
|
||||
|
||||
// Wraps ort_api->GetOnnxTypeFromTypeInfo.
|
||||
OrtStatus *GetONNXTypeFromTypeInfo(OrtTypeInfo *info, enum ONNXType *out);
|
||||
|
||||
// Wraps ort_api->FreeTypeInfo.
|
||||
void ReleaseTypeInfo(OrtTypeInfo *o);
|
||||
|
||||
// Wraps ort_spi->SessionGetModelMetadata.
|
||||
OrtStatus *SessionGetModelMetadata(OrtSession *s, OrtModelMetadata **out);
|
||||
|
||||
// Wraps ort_api->ReleaseModelMetadata.
|
||||
void ReleaseModelMetadata(OrtModelMetadata *m);
|
||||
|
||||
// Wraps ort_api->ModelMetadataGetProducerName, using the default allocator.
|
||||
OrtStatus *ModelMetadataGetProducerName(OrtModelMetadata *m, char **name);
|
||||
|
||||
// Wraps ort_api->ModelMetadataGetGraphName, using the default allocator.
|
||||
OrtStatus *ModelMetadataGetGraphName(OrtModelMetadata *m, char **name);
|
||||
|
||||
// Wraps ort_api->ModelMetadataGetDomain, using the default allocator.
|
||||
OrtStatus *ModelMetadataGetDomain(OrtModelMetadata *m, char **domain);
|
||||
|
||||
// Wraps ort_api->ModelMetadataGetDescription, using the default allocator.
|
||||
OrtStatus *ModelMetadataGetDescription(OrtModelMetadata *m, char **desc);
|
||||
|
||||
// Wraps ort_api->ModelMetadataLookupCustomMetadataMap, using the default
|
||||
// allocator.
|
||||
OrtStatus *ModelMetadataLookupCustomMetadataMap(OrtModelMetadata *m, char *key,
|
||||
char **value);
|
||||
|
||||
// Wraps ort_api->ModelMetadataGetCustomMetadataMapKeys, using the default
|
||||
// allocator.
|
||||
OrtStatus *ModelMetadataGetCustomMetadataMapKeys(OrtModelMetadata *m,
|
||||
char ***keys, int64_t *num_keys);
|
||||
|
||||
// Wraps ort_api->ModelMetadataGetVersion.
|
||||
OrtStatus *ModelMetadataGetVersion(OrtModelMetadata *m, int64_t *version);
|
||||
|
||||
// Wraps ort_api->GetValue. Uses the default allocator.
|
||||
OrtStatus *GetValue(OrtValue *container, int index, OrtValue **dst);
|
||||
|
||||
// Wraps ort_api->GetValueType.
|
||||
OrtStatus *GetValueType(OrtValue *v, enum ONNXType *out);
|
||||
|
||||
// Wraps ort_api->GetValueCount.
|
||||
OrtStatus *GetValueCount(OrtValue *v, size_t *out);
|
||||
|
||||
// Wraps ort_api->CreateValue to create a map or a sequence.
|
||||
OrtStatus *CreateOrtValue(OrtValue **in, size_t num_values,
|
||||
enum ONNXType value_type, OrtValue **out);
|
||||
|
||||
// TRAINING API WRAPPER
|
||||
|
||||
void SetTrainingApi();
|
||||
|
||||
// Checks if training api is supported.
|
||||
int IsTrainingApiSupported();
|
||||
|
||||
// Wraps ort_training_api->CreateSessionFromBuffer.
|
||||
// Creates and ORT checkpoint from the checkpoint data.
|
||||
OrtStatus *CreateCheckpoint(void *checkpoint_data,
|
||||
size_t checkpoint_data_length, OrtCheckpointState **out);
|
||||
|
||||
// Wraps ort_training_api->CreateTrainingSessionFromBuffer. Creates an ORT
|
||||
// training session using the given models and checkpoint. The given options
|
||||
// pointer may be NULL; if it is, then we'll use default options.
|
||||
OrtStatus *CreateTrainingSessionFromBuffer(OrtCheckpointState *checkpoint_state,
|
||||
void *training_model_data, size_t training_model_data_length,
|
||||
void *eval_model_data, size_t eval_model_data_length,
|
||||
void *optim_model_data, size_t optim_model_data_length,
|
||||
OrtEnv *env, OrtTrainingSession **out, OrtSessionOptions *options);
|
||||
|
||||
// Wraps ort_training_api->CreateTrainingSession.
|
||||
// Currently this is the only way to create a training session that is able to
|
||||
// export the final trained model to disk.
|
||||
OrtStatus *CreateTrainingSessionFromPaths(OrtCheckpointState *checkpoint_state,
|
||||
char *training_model_path, char *eval_model_path, char *optim_model_path,
|
||||
OrtEnv *env, OrtTrainingSession **out, OrtSessionOptions *options);
|
||||
|
||||
// Wraps ort_training_api->TrainingSessionGetTrainingModelInputCount
|
||||
// and ort_training_api->TrainingSessionGetEvalgModelInputCount.
|
||||
OrtStatus *TrainingSessionGetInputCount(OrtTrainingSession *training_session, size_t *result_training, size_t *result_eval);
|
||||
|
||||
// Wraps ort_training_api->TrainingSessionGetTrainingModelOutputCounet
|
||||
// and ort_training_api->TrainingSessionGetEvalgModelOutputCount.
|
||||
OrtStatus *TrainingSessionGetOutputCount(OrtTrainingSession *training_session, size_t *result_training, size_t *result_eval);
|
||||
|
||||
// Wraps ort_training_api->TrainingSessionGetTrainingModelInputName.
|
||||
OrtStatus *TrainingSessionGetTrainingInputName(OrtTrainingSession *training_session, size_t i, char **name);
|
||||
|
||||
// Wraps ort_training_api->TrainingSessionGetEvalModelInputName.
|
||||
OrtStatus *TrainingSessionGetEvalInputName(OrtTrainingSession *training_session, size_t i, char **name);
|
||||
|
||||
// Wraps ort_training_api->TrainingSessionGetTrainingModelOutputName.
|
||||
OrtStatus *TrainingSessionGetTrainingOutputName(OrtTrainingSession *training_session, size_t i, char **name);
|
||||
|
||||
// Wraps ort_training_api->TrainingSessionGetEvalModelOutputName.
|
||||
OrtStatus *TrainingSessionGetEvalOutputName(OrtTrainingSession *training_session, size_t i, char **name);
|
||||
|
||||
// Wraps ort_training_api->TrainStep.
|
||||
OrtStatus *TrainStep(OrtTrainingSession *training_session, size_t inputs_len, OrtValue **inputs, size_t output_len, OrtValue **outputs);
|
||||
|
||||
// Wraps ort_training_api->OptimizerStep.
|
||||
OrtStatus *OptimizerStep(OrtTrainingSession *training_session);
|
||||
|
||||
// Wraps ort_training_api->LazyResetGrad.
|
||||
OrtStatus *LazyResetGrad(OrtTrainingSession *training_session);
|
||||
|
||||
// Wraps ort_training_api->SaveCheckpoint.
|
||||
OrtStatus *SaveCheckpoint(OrtCheckpointState *checkpoint, char *path, size_t include_optimizer);
|
||||
|
||||
// Wraps ort_training_api->ExportModel.
|
||||
OrtStatus *ExportModel(OrtTrainingSession *training_session, char *path, size_t outputs_len, char **output_names);
|
||||
|
||||
// Wraps ort_training_api->ReleaseTrainingSession.
|
||||
void ReleaseOrtTrainingSession(OrtTrainingSession *session);
|
||||
|
||||
// Wraps ort_training_api->ReleaseCheckpointState.
|
||||
void ReleaseCheckpointState(OrtCheckpointState *checkpoint);
|
||||
|
||||
#ifdef __cplusplus
|
||||
} // extern "C"
|
||||
#endif
|
||||
#endif // ONNXRUNTIME_WRAPPER_H
|
||||
+109
@@ -0,0 +1,109 @@
|
||||
//go:build !windows
|
||||
|
||||
package onnxruntime_go
|
||||
|
||||
import (
|
||||
"fmt"
|
||||
"runtime"
|
||||
"unsafe"
|
||||
)
|
||||
|
||||
/*
|
||||
#cgo LDFLAGS: -ldl
|
||||
|
||||
#include <dlfcn.h>
|
||||
#include "onnxruntime_wrapper.h"
|
||||
|
||||
typedef OrtApiBase* (*GetOrtApiBaseFunction)(void);
|
||||
|
||||
// Since Go can't call C function pointers directly, we just use this helper
|
||||
// when calling GetApiBase
|
||||
OrtApiBase *CallGetAPIBaseFunction(void *fn) {
|
||||
OrtApiBase *to_return = ((GetOrtApiBaseFunction) fn)();
|
||||
return to_return;
|
||||
}
|
||||
*/
|
||||
import "C"
|
||||
|
||||
// This file includes the code for loading the onnxruntime and setting up the
|
||||
// environment on non-Windows systems. For now, it has been tested on Linux and
|
||||
// arm64 OSX.
|
||||
|
||||
// This will contain the handle to the onnxruntime shared library if it has
|
||||
// been loaded successfully.
|
||||
var libraryHandle unsafe.Pointer
|
||||
|
||||
func platformCleanup() error {
|
||||
v, e := C.dlclose(libraryHandle)
|
||||
if v != 0 {
|
||||
return fmt.Errorf("Error closing the library: %w", e)
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// Should only be called on Apple systems; looks up the CoreML provider
|
||||
// function which should only be exported on apple onnxruntime dylib files.
|
||||
func setAppendCoreMLFunctionPointer(libraryHandle unsafe.Pointer) error {
|
||||
// This function name must match the name in coreml_provider_factory.h,
|
||||
// which is provided in the onnxruntime release's include/ directory on for
|
||||
// Apple platforms.
|
||||
fnName := "OrtSessionOptionsAppendExecutionProvider_CoreML"
|
||||
cFunctionName := C.CString(fnName)
|
||||
defer C.free(unsafe.Pointer(cFunctionName))
|
||||
appendCoreMLProviderProc := C.dlsym(libraryHandle, cFunctionName)
|
||||
if appendCoreMLProviderProc == nil {
|
||||
msg := C.GoString(C.dlerror())
|
||||
return fmt.Errorf("Error looking up %s: %s", fnName, msg)
|
||||
}
|
||||
C.SetCoreMLProviderFunctionPointer(appendCoreMLProviderProc)
|
||||
return nil
|
||||
}
|
||||
|
||||
func platformInitializeEnvironment() error {
|
||||
if onnxSharedLibraryPath == "" {
|
||||
onnxSharedLibraryPath = "onnxruntime.so"
|
||||
}
|
||||
cName := C.CString(onnxSharedLibraryPath)
|
||||
defer C.free(unsafe.Pointer(cName))
|
||||
handle := C.dlopen(cName, C.RTLD_LAZY)
|
||||
if handle == nil {
|
||||
msg := C.GoString(C.dlerror())
|
||||
return fmt.Errorf("Error loading ONNX shared library \"%s\": %s",
|
||||
onnxSharedLibraryPath, msg)
|
||||
}
|
||||
cFunctionName := C.CString("OrtGetApiBase")
|
||||
defer C.free(unsafe.Pointer(cFunctionName))
|
||||
getAPIBaseProc := C.dlsym(handle, cFunctionName)
|
||||
if getAPIBaseProc == nil {
|
||||
C.dlclose(handle)
|
||||
msg := C.GoString(C.dlerror())
|
||||
return fmt.Errorf("Error looking up OrtGetApiBase in \"%s\": %s",
|
||||
onnxSharedLibraryPath, msg)
|
||||
}
|
||||
ortAPIBase := C.CallGetAPIBaseFunction(getAPIBaseProc)
|
||||
tmp := C.SetAPIFromBase((*C.OrtApiBase)(unsafe.Pointer(ortAPIBase)))
|
||||
if tmp != 0 {
|
||||
C.dlclose(handle)
|
||||
return fmt.Errorf("Error setting ORT API base: %d", tmp)
|
||||
}
|
||||
if (runtime.GOOS == "darwin") || (runtime.GOOS == "ios") {
|
||||
setAppendCoreMLFunctionPointer(handle)
|
||||
// We'll silently ignore potential errors returned by
|
||||
// setAppendCoreMLFunctionPointer (for now at least). Even though we're
|
||||
// on Apple hardware, it's possible that the user will have compiled
|
||||
// the onnxruntime library from source without CoreML support.
|
||||
// A failure here will only leave the coreml function pointer as NULL
|
||||
// in our C code, which will be detected and result in an error at
|
||||
// runtime.
|
||||
}
|
||||
libraryHandle = handle
|
||||
return nil
|
||||
}
|
||||
|
||||
// Converts the given path to an ORTCHAR_T string, pointed to by a *C.char. The
|
||||
// returned string must be freed using C.free when no longer needed. This
|
||||
// wrapper is used for source compatibility with onnxruntime API functions
|
||||
// requiring paths, which must be UTF-16 on Windows but UTF-8 elsewhere.
|
||||
func createOrtCharString(str string) (*C.char, error) {
|
||||
return C.CString(str), nil
|
||||
}
|
||||
+107
@@ -0,0 +1,107 @@
|
||||
//go:build windows
|
||||
|
||||
package onnxruntime_go
|
||||
|
||||
// This file includes the Windows-specific code for loading the onnxruntime
|
||||
// library and setting up the environment.
|
||||
|
||||
import (
|
||||
"fmt"
|
||||
"syscall"
|
||||
"unicode/utf16"
|
||||
"unicode/utf8"
|
||||
"unsafe"
|
||||
)
|
||||
|
||||
// #include "onnxruntime_wrapper.h"
|
||||
import "C"
|
||||
|
||||
// This will contain the handle to the onnxruntime dll if it has been loaded
|
||||
// successfully.
|
||||
var libraryHandle syscall.Handle
|
||||
|
||||
func platformCleanup() error {
|
||||
e := syscall.FreeLibrary(libraryHandle)
|
||||
libraryHandle = 0
|
||||
return e
|
||||
}
|
||||
|
||||
func platformInitializeEnvironment() error {
|
||||
if onnxSharedLibraryPath == "" {
|
||||
onnxSharedLibraryPath = "onnxruntime.dll"
|
||||
}
|
||||
handle, e := syscall.LoadLibrary(onnxSharedLibraryPath)
|
||||
if e != nil {
|
||||
return fmt.Errorf("Error loading ONNX shared library \"%s\": %w",
|
||||
onnxSharedLibraryPath, e)
|
||||
}
|
||||
getApiBaseProc, e := syscall.GetProcAddress(handle, "OrtGetApiBase")
|
||||
if e != nil {
|
||||
syscall.FreeLibrary(handle)
|
||||
return fmt.Errorf("Error finding OrtGetApiBase function in %s: %w",
|
||||
onnxSharedLibraryPath, e)
|
||||
}
|
||||
ortApiBase, _, e := syscall.SyscallN(uintptr(getApiBaseProc), 0)
|
||||
if ortApiBase == 0 {
|
||||
syscall.FreeLibrary(handle)
|
||||
if e != nil {
|
||||
return fmt.Errorf("Error calling OrtGetApiBase: %w", e)
|
||||
} else {
|
||||
return fmt.Errorf("Error calling OrtGetApiBase")
|
||||
}
|
||||
}
|
||||
tmp := C.SetAPIFromBase((*C.OrtApiBase)(unsafe.Pointer(ortApiBase)))
|
||||
if tmp != 0 {
|
||||
syscall.FreeLibrary(handle)
|
||||
return fmt.Errorf("Error setting ORT API base: %d", tmp)
|
||||
}
|
||||
|
||||
// we do not initialize the training API on windows (see setup_env.go)
|
||||
// because currently we cannot support the conversion from UTF-8 to wide
|
||||
// character. See https://github.com/yalue/onnxruntime_go/pull/56.
|
||||
|
||||
libraryHandle = handle
|
||||
return nil
|
||||
}
|
||||
|
||||
// Converts the given string to a UTF-16 string, pointed to by a raw
|
||||
// *C.char. Note that we actually keep ORTCHAR_T defined to char even
|
||||
// on Windows, so do _not_ index into this string from Cgo code and expect to
|
||||
// get correct characters! Instead, this should only be used to obtain pointers
|
||||
// that are passed to onnxruntime windows DLL functions expecting ORTCHAR_T*
|
||||
// args. This is required because we undefine _WIN32 for cgo compatibility when
|
||||
// including onnxruntime_c_api.h, but still interact with a DLL that was
|
||||
// compiled assuming _WIN32 was defined.
|
||||
//
|
||||
// The pointer returned by this function must still be freed using C.free when
|
||||
// no longer needed. This will return an error if the given string contains
|
||||
// non-UTF8 characters.
|
||||
func createOrtCharString(str string) (*C.char, error) {
|
||||
src := []uint8(str)
|
||||
// Assumed common case: the utf16 buffer contains one uint16 per utf8 byte
|
||||
// plus one more for the required null terminator in the C buffer.
|
||||
dst := make([]uint16, 0, len(src)+1)
|
||||
// Convert UTF-8 to UTF-16 by reading each subsequent rune from src and
|
||||
// appending it as UTF-16 to dst.
|
||||
for len(src) > 0 {
|
||||
r, size := utf8.DecodeRune(src)
|
||||
if r == utf8.RuneError {
|
||||
return nil, fmt.Errorf("Invalid UTF-8 rune found in \"%s\"", str)
|
||||
}
|
||||
src = src[size:]
|
||||
dst = utf16.AppendRune(dst, r)
|
||||
}
|
||||
// Make sure dst contains the null terminator. Additionally this will cause
|
||||
// us to return an empty string if the original string was empty.
|
||||
dst = append(dst, 0)
|
||||
|
||||
// Finally, we need to copy dst into a C array for compatibility with
|
||||
// C.CString.
|
||||
toReturn := C.calloc(C.size_t(len(dst)), 2)
|
||||
if toReturn == nil {
|
||||
return nil, fmt.Errorf("Error allocating buffer for the utf16 string")
|
||||
}
|
||||
C.memcpy(toReturn, unsafe.Pointer(&(dst[0])), C.size_t(len(dst))*2)
|
||||
|
||||
return (*C.char)(toReturn), nil
|
||||
}
|
||||
+54
@@ -0,0 +1,54 @@
|
||||
package onnxruntime_go
|
||||
|
||||
// This file contains definitions for the generic tensor data types we support.
|
||||
|
||||
// #include "onnxruntime_wrapper.h"
|
||||
import "C"
|
||||
|
||||
import (
|
||||
"reflect"
|
||||
)
|
||||
|
||||
type FloatData interface {
|
||||
~float32 | ~float64
|
||||
}
|
||||
|
||||
type IntData interface {
|
||||
~int8 | ~uint8 | ~int16 | ~uint16 | ~int32 | ~uint32 | ~int64 | ~uint64
|
||||
}
|
||||
|
||||
// This is used as a type constraint for the generic Tensor type.
|
||||
type TensorData interface {
|
||||
FloatData | IntData
|
||||
}
|
||||
|
||||
// Returns the ONNX enum value used to indicate TensorData type T.
|
||||
func GetTensorElementDataType[T TensorData]() C.ONNXTensorElementDataType {
|
||||
// Sadly, we can't do type assertions to get underlying types, so we need
|
||||
// to use reflect here instead.
|
||||
var v T
|
||||
kind := reflect.ValueOf(v).Kind()
|
||||
switch kind {
|
||||
case reflect.Float64:
|
||||
return C.ONNX_TENSOR_ELEMENT_DATA_TYPE_DOUBLE
|
||||
case reflect.Float32:
|
||||
return C.ONNX_TENSOR_ELEMENT_DATA_TYPE_FLOAT
|
||||
case reflect.Int8:
|
||||
return C.ONNX_TENSOR_ELEMENT_DATA_TYPE_INT8
|
||||
case reflect.Uint8:
|
||||
return C.ONNX_TENSOR_ELEMENT_DATA_TYPE_UINT8
|
||||
case reflect.Int16:
|
||||
return C.ONNX_TENSOR_ELEMENT_DATA_TYPE_INT16
|
||||
case reflect.Uint16:
|
||||
return C.ONNX_TENSOR_ELEMENT_DATA_TYPE_UINT16
|
||||
case reflect.Int32:
|
||||
return C.ONNX_TENSOR_ELEMENT_DATA_TYPE_INT32
|
||||
case reflect.Uint32:
|
||||
return C.ONNX_TENSOR_ELEMENT_DATA_TYPE_UINT32
|
||||
case reflect.Int64:
|
||||
return C.ONNX_TENSOR_ELEMENT_DATA_TYPE_INT64
|
||||
case reflect.Uint64:
|
||||
return C.ONNX_TENSOR_ELEMENT_DATA_TYPE_UINT64
|
||||
}
|
||||
return C.ONNX_TENSOR_ELEMENT_DATA_TYPE_UNDEFINED
|
||||
}
|
||||
Reference in New Issue
Block a user