All of this is absurdly complicated. Exactly what I would expect from a new student who doesn't know what they're doing and has no one to teach them how do you engineering in a systematic manner. I don't mean this as an insult. I teach this stuff and have seen it hundreds of times.
You should look for "post training static quantization" also called . There are countless ways to quantize. This will quantize both the weights and the activations after training.
You're doing this on hard mode for no reason. This is typical and something I often need to break people out of. Optimizing for performance by doing custom things in Jax when you're a beginner is a terrible path to take.
Performance is not your problem. You're training a trivial network that would have run on a CPU 20 years ago.
There's no clear direction here, just trying complicated stuff in no logical order with no learning or dependencies between steps. You need to treat these problems as scientific experiments. What do I do to learn more about my domain, what do I change depending on the answer I get, etc. Not, now it's time to try something else random like jax.
Worse. You need to learn the key lesson in this space. Credit assignment for problems is extremely hard. If something isn't working why isn't it? Because of a bug? A hopeless problem? Using a crappy optimizer? Etc. That's why you should start in a framework that works and escape it later if you want.
Here's a simple plan to do this:
First forget about quantization. Use pytorch. Implement your trivial network in 5 lines. Train it with Adam. Make sure it works. Make sure your problem is solveable with the data that you have and the network you've chosen and your activation functions and the loss and the optimizer (use Adam, forget about this doing stuff by hand for now).
> Unless I had an expert guide who was absolutely sure it’d be straightforward (email me!), I’d avoid high-level frameworks like TensorFlow and PyTorch and instead implement the quantization-aware training myself.
This is exactly backwards. Unless you have an expert never implement anything yourself. If you don't have one, rely on what already exists. Because you can logically narrow down the options for what works and what's wrong. If you do it yourself you're always lost.
Once you have that working start backing off. Slowly change the working network into what you need. Step by step. At every step write down why you think your change is good and what you would do if it isn't. Then look at the results.
Forget about microflow-rs or whatever. Train with pytorch, export to onnx, generate c code for your onnx for inference.
I kind of see your point, but only in the context of working on time-sensitive task which others rely upon. But if it is hobby/educational project, what is wrong doing things by yourself? And resort to decomposing existing solution if you can't figure out why yours is not working?
There's nothing better for understanding something rather than trying to do that "something" from scratch yourself.
I think the point is that OP is learning things about a wide variety of topics that aren't really relevant to their stated goal, i.e. solving the sensor/state inference problem.
Which, as you say, can be valuable! There's nothing wrong with that. But the more complexity you add the less likely you are to actually solve the problem (all else being equal, some problems are just inherently complex).
You should look for "post training static quantization" also called . There are countless ways to quantize. This will quantize both the weights and the activations after training.
You're doing this on hard mode for no reason. This is typical and something I often need to break people out of. Optimizing for performance by doing custom things in Jax when you're a beginner is a terrible path to take.
Performance is not your problem. You're training a trivial network that would have run on a CPU 20 years ago.
There's no clear direction here, just trying complicated stuff in no logical order with no learning or dependencies between steps. You need to treat these problems as scientific experiments. What do I do to learn more about my domain, what do I change depending on the answer I get, etc. Not, now it's time to try something else random like jax.
Worse. You need to learn the key lesson in this space. Credit assignment for problems is extremely hard. If something isn't working why isn't it? Because of a bug? A hopeless problem? Using a crappy optimizer? Etc. That's why you should start in a framework that works and escape it later if you want.
Here's a simple plan to do this:
First forget about quantization. Use pytorch. Implement your trivial network in 5 lines. Train it with Adam. Make sure it works. Make sure your problem is solveable with the data that you have and the network you've chosen and your activation functions and the loss and the optimizer (use Adam, forget about this doing stuff by hand for now).
> Unless I had an expert guide who was absolutely sure it’d be straightforward (email me!), I’d avoid high-level frameworks like TensorFlow and PyTorch and instead implement the quantization-aware training myself.
This is exactly backwards. Unless you have an expert never implement anything yourself. If you don't have one, rely on what already exists. Because you can logically narrow down the options for what works and what's wrong. If you do it yourself you're always lost.
Once you have that working start backing off. Slowly change the working network into what you need. Step by step. At every step write down why you think your change is good and what you would do if it isn't. Then look at the results.
Forget about microflow-rs or whatever. Train with pytorch, export to onnx, generate c code for your onnx for inference.
Read the pytorch guide on PTSQ and use it.