Open josephrocca opened 2 years ago
+1
+1 When using https://tfhub.dev/tulasiram58827/lite-model/keras-ocr/dr/2 with tfjs in browser. It throws this error when loading the model.
+1
+1
My best guess is JAX-MD requires a dynamically loaded library, or JAX is trying to dynamically load JAX-MD in the browser (perhaps someone more familiar with the library can confirm this). This isn't supported yet, but we're exploring adding dynamically loaded TFLite delegates to tfjs-tflite. If that works, we may be able to support this as well.
Hey @mattsoulanille, this is just a normal tflite file - you can download it here and see it in Netron (see image below). So unless I misunderstand you, there is no JAX-MD 'library' that's being loaded/used in the browser. I just used some functions from it and then converted my model from JAX to tflite. So there shouldn't be any JAX-MD-specific stuff happening at runtime - it's just a regular tflite model at this point. Please forgive any misunderstandings I might have here though!
Note that the full code to reproduce is just:
<script src="https://cdn.jsdelivr.net/npm/@tensorflow/tfjs-core@3.14.0/dist/tf-core.js"></script>
<script src="https://cdn.jsdelivr.net/npm/@tensorflow/tfjs-backend-cpu@3.14.0/dist/tf-backend-cpu.js"></script>
<script src="https://cdn.jsdelivr.net/npm/@tensorflow/tfjs-tflite@0.0.1-alpha.8/dist/tf-tflite.js"></script>
<script type="module">
const initFn = await tflite.loadTFLiteModel('https://josephrocca.github.io/jax-md-web-test/minimal/init_fn_raw.tflite');
</script>
Thanks for the details. I've taken a deeper look, and I'm pretty sure this is caused by using TF select ops in the TensorFlow to TFLite conversion. Select ops require linking ops from TensorFlow in the TFLite binary when compiling TFLite, but the @tensorflow/tfjs-tflite
wasm binary does not include these ops. I think the dlopen
error you're seeing is caused by the TFLite interpreter seeing a TF op that it does not include and trying to dynamically load TF with dlopen
instead (I may be wrong about this).
Right now, the build infrastructure for @tensorflow/tfjs-tflite
's wasm binary is internal to Google, but we want to open-source it to make these kinds of problems easier for OSS users to work around. If it were open-source right now, I'd say you can try building a custom version of the wasm file that supports the ops you need. Another approach (once this is open-sourced) might be to build TFLite with dlopen
support and provide a minimal TF library that implements the ops you need. However, at the moment, I don't have a good solution. Sorry.
I'll keep this thread updated as we open-source the tfjs-tflite build infrastructure.
Thanks @mattsoulanille, I think you're right! Great to hear you're open sourcing the tfjs-tflite build infrastructure. It would be really awesome if, as part of that open-sourcing process, you made it really easy for users to compile their own tflite runtime that includes the specific ops they need, or even just provided a runtime that includes all the tf select ops as a "bad" fallback, that at least allows users to get their model working so they're unblocked, and can come back to optimise the size of the runtime later. I talk more about this here.
I've been using ONNX Runtime Web more and more often lately because so many models have TF Select ops, or unsupported ops in tfjs - it would be really good to have a "fallback" runtime on the tfjs side of things that "just works", even if it's a bit bloated. That said, if you end up making it super easy for users to compile their own tfjs-tflite runtime, then that probably wouldn't be necessary.
Is there a fix for this? I have the same error..
Also, how can I build Tensorflow Lite with dlopen
support?
Thanks!
Hi @compromyse. We're still working on open-sourcing the tfjs-tflite build infrastructure, but I'll be sure to update this thread when it's ready.
Hi @compromyse. We're still working on open-sourcing the tfjs-tflite build infrastructure, but I'll be sure to update this thread when it's ready.
Sure, thanks!
Can you compile a version of tflite that enables dlopen
? (I'm not sure why it has to be open source to do that).
+1 to this issue, unable to load TFLite model with select ops..
is it possible to compile a wasm with select ops and host it publicly before open sourcing it?
Sorry, any update so far?
System information
I'm trying to get the JAX-MD molecular dynamics library working in the browser, and I've run into this error after converting a simple model from JAX to tf with jax2tf, and then from tf to tflite with
tf.lite.TFLiteConverter.from_saved_model
. Here's the notebook that does the conversions.When running the model in the browser I get this error:
Standalone code to reproduce the issue
Additional notes I noticed that there's a StackOverflow question about this (same error as above when trying to use tfjs-tflite with some models) with about 100 views as of writing so it seems that there are others bumping into this too. I wasn't able to find an existing issue in this repo for it though.