Add support for tf.SavedModel format about neuropod HOT 7 CLOSED

I spent a bunch of time on this over the last 9 hours and now have an initial SavedModel implementation that passes tests from both Python and C++ (with TF 2.2)!

It was fairly complex to get working, but I think there's a reasonable path forward.

I'll post more context in this thread later this week.

from neuropod.

Hey everyone,

Here's the context I mentioned above:

Context

Why doesn't Neuropod already have support for SavedModels?

The initial TF implementation for Neuropod was built when ATG was using TF (1.x) frozen graphs directly. All of our models required only frozen graphs to run inference so SavedModel support wasn't a pressing need.

I remember spending a decent amount of time exploring SavedModel support in the past, but there was significant overhead to getting it working well in Neuropod (compatibility with multiple TF versions, relatively complex integration, making it easy for users to export models without having to provide the same set of information multiple times, etc).

Our existing solution fit our needs (and the needs of all of our users as far as we were aware) so the implementation and design effort wasn't worth it at the time. ATG also started focusing on PyTorch for newer models so we didn't have a compelling reason to make major changes to something that wasn't broken.

Why not use the built-in LoadSavedModel in Tensorflow?

TF doesn't provide a good way to run graphs on different devices in the same process.

For example, we can't use GPUOptions::visible_device_list as it is a per process setting

From: tensorflow/tensorflow#18861 (comment)

Unfortunately, though visible_deivces_list is included in ConfigProto, it is
actually a per-process setting. In fact, this is true of almost all options inside
the GPUOptions protocol buffer.

The issue linked above goes through several other alternatives that don't work either.

Instead, we have to set the device of the graph by moving each node in the GraphDef to the target device before creating a session.

LoadSavedModel creates a session and loads a GraphDef internally so we can't use it directly.

Why not convert SavedModels to Frozen Graphs?

I'm not super confident about this, so feel free to correct me if you have an in-depth knowledge of SavedModels:

Based on the research I've done, I think it is possible to convert all saved models to frozen graphs by loading the SavedModel into a session, running all init ops and dumping the resulting graph. I believe this should make it independent of any asset files on disk (as I believe that all assets are loaded into tensors during init).

I think a more important question though, is whether we want to do this or not. It requires changing the user's model instead of just wrapping it and also may not be compatible with all SavedModel features. I'm not concretely aware of SavedModel features that aren't supported by frozen graphs at the moment, but I don't believe there's a compatibility guarantee.

My implementation

The approach I took in my implementation was to pull in the SavedModel loading code from TF 2.4 into the Neuropod TF backend. There was one major modification and some minor modifications (to remove unneeded code and change header paths).

The major modification was to do the device rewriting operation mentioned earlier in my comment before loading the GraphDef.

This should make Neuropod compatible with all SavedModel features as we're effectively using the actual SavedModel loading code from TensorFlow.

The rest of this week is quite busy for me, but I'll try to get the PR up early next week. Still have a bit of cleanup work to do before putting it up.

from neuropod.

I agree we should add support for TF SavedModels. If anyone's willing to help with implementing this, please comment here.

Also, to clarify, Neuropod does support packaging TF .pb files. create_tensorflow_neuropod lets you pass in a GraphDef or a frozen graph (a .pb file you've already exported)

from neuropod.

Vivek,

Do you use it to convert to frozen graph or use it to initialize?

FYI: I did conversion once and when looked for generic solution found this code https://github.com/ogemarques/models/blob/master/research/tensorrt/tensorrt.py#L163

from neuropod.

Hey Vivek, this is great, thank you! Looking forward for it.

from neuropod.

Can I try and work on this?

from neuropod.

Can I try and work on this?

Hi @ArkaprabhaChakraborty! Thanks for the offer!

There's already a PR up for this (#487), but there are other issues that are probably better for a first contribution. Maybe #138 or #351?

If you're interested in contributing, starting with one or two smaller issues (like the ones mentioned above) provides a good introduction to the codebase, build system, running tests locally, etc. After that, it becomes easier to dig into more complex issues and features like this one.

There's also always room for improvement for documentation if you're interested in that.

Feel free to comment on either of those issues if you're interested or have questions. Also, take a look at the contribution guide if you haven't already.

Thanks again for the offer to help!

from neuropod.