Running pre-trained ML models in Godot
Techie Kho
Posted on January 18, 2024
So I have been developing this GDExtension called iree.gd. It is mission to embed IREE, another cool project that compiles and runs ML models, into Godot. It took me quite a while, but finally It has reached alpha.
Hope you guys could check it out the sample.
I can't upload video to here, but you still can watch a small clips of iree.gd doing its magic here
Overview
Currently, it is still having a very thin abstraction layer, and some manual stuff need to be done in prior.
Preparing the ML model
Before doing anything, you'll need to download your trained ML model from the internet. You could find some from Kaggle. Then, you'll need to preprocess them with python via this detailed documentation. You'll need to set the HAL backend following this table depending on your targeted platform. As you generate the .vmfb
byte code, you are ready to proceed to the next step.
Use the ML model
.vmfb
files are automatically imported into the project as you drag and drop 'em into it. You could then use them in your script by loading them.
var model : IREEModule = preload("res://your_model.vmfb")
@export var model_2 : IREEModule # Drag and drop in editor
Then, you'll need to prepare your inputs with IREETensor
, then execute the ML model.
var input := IREETensor.from_bytes(image.get_data(), [1, 50, 50, 3]) # Remember to consider the input type.
var outputs := module.run_module("module.main", [input])
for output in outputs:
pass # Do something with the `output`.
The input
and function name vary models to models. You could inspect the .vmfb
file using iree-dump-module
from the IREE packages you installed during preparing the ML model.
Here is one section of the dump of the ML model used:
...
Exported Functions:
[ 0] main(!vm.ref<?>) -> (!vm.ref<?>)
iree.abi.declaration: sync func @main(%input0: tensor<1x50x50x3xf32> {ml_program.identifier = "input_0"}) -> (%output0: tensor<1x200x200x3xf32> {ml_program.identifier = "Identity"})
[ 1] __init() -> ()
...
You could kind of guess it, it has main
function that takes a tensor of float32
(1x50x50x3xf32
) and outputs another tensor of float32
(1x200x200x3xf32
).
Sample
There is a readily available sample of esrgan for you to try in the release page.
The sample only work on windows and linux.
After download and open iree-gd-sample-*.zip
with Godot 4.2+, you will need to do these:
- set
res://scenes/main/main.tscn
to be loaded on start. - Drag and drop appropriate byte code in
res://bytecodes/esrgan/esrgan.*.vmfb
intoUI/TextureRect
's module export. - Run it.
Afterword
Hopefully, in the future, Godot developer could utilise this piece of technology and create something extraordinary and magical. Thank you for reading this.
Posted on January 18, 2024
Join Our Newsletter. No Spam, Only the good stuff.
Sign up to receive the latest update from our blog.