I've been playing around with Keras's built-in visualizations tool for a while and I feel like it could be improved a lot. The current visualization system doesn't provide much information "at a glance" (meaning it takes quite a bit of analysis to understand what's going on) and isn't super aesthetically pleasing.

I'm opening this issue to see if an improved visualization system is something that the Keras team would be open to. I'd be happy to make small improvements to the existing system, but I feel like it would be better to add a new visualization system (or submodule) for highly customizable visual architecture generation. I could bring some (or most) of the key features from visualkeras, a repository I maintain, into the src/visualization module.

Since we'd isolate the dependencies of the new features from Keras as a whole (like what was done for the last system), we wouldn't be adding any bloat through more dependencies.

Below are a few examples of what the visualizations I'd add may look like

Image

Image

Image

Image

Image

Comment From: mattdangerw

I think we'd be most interested in visualizations that could work for generic functional models (not purely sequential models with a one input and one output constraint). Even for convnets you'll usually have residuals, so most real world model will not be purely sequential.

What would the "graph view" you are proposing look like, and how would it compare to plot_model?

Comment From: Soontosh

Thanks for the response. The visualkeras features I was planning to import were the functions layered_view and graph_view. They process a tf.keras.Model (or keras.Model) into a PIL image showing each layer as a scaled 3D block stack or as a graph (see the orange graph visualization attached). There are additionally many different customization options (e.g., converting 3D block stack into 2D).

So it takes in any Keras Model instance plus a handful of optional kwargs and returns a Pillow image. The only extra dependencies are Pillow and aggdraw which would be optional dependencies just for this feature (similar to how the visualization component currently works) so that we do not add any extra “bloat” or anything.

We do have limited support for functional models, but it currently works mainly with simple functional models that are single stream. To make this work better with Functional models, I’d traverse model._nodes_by_depth to build a list, draw skip/merge edges as curved connectors, seed the traversal from every Model.inputs for multi-input/output architectures, and optionally fold the repeated subgraphs into a single block with an “xN” label. We’d have to make sure each different input stream has proper dynamically generated spacing so that it’s not cluttered (we could have user control to some extent here). This is just a rough overview I’ve come up with and it will probably require changes in implementation. I am unsure if we could get that working properly with 3D views, but this is definitely doable with a graph-based view.

Please let me know what you think would be the best thing for me to implement or modify (and how I should do so) to improve the Keras visualization system.