CNTK 102.5: Visualizing CNTK models with Graphviz and D3 in F# for Jupyter

CNTK 102.5: Visualizing CNTK models with Graphviz and D3 in F# for Jupyter

 

Retroactive changes


  • Notebooks from now on will be on their own folders, along will their corresponding .fsx

This makes is so much easier to keep track of dependencies, and we also resolve the issue of IfCntk overriding the local main.group.fsx path. Old links will still work.

 

Intro

The Python API offers a simple way to display a visual graph of your model via the well known Graphviz utility, that sadly hasn't made it to the .NET Cntk API so far.

It's very useful to be able to have an idea at a glance of what is going on inside your model, especially if you are not the one who originally created it. There have been other implementations towards that end, and in F# no less, such as Faisal Waris' FsCNTKTools that uses Microsoft's Automatic Graph Layout module for visualization.

Here we are going to use the a d3js implementation of Graphviz called d3-graphviz that although not a 100% port, is pretty adequate for our needs.

This configuration also presents us with certain advantages: first, you absolutely don't need to have any Graphviz binaries installed to display graphs, which has been known to cause issues even in IPython, and second, since the implementation is built on d3js we will have a great deal of options to add interactivity to the graphs, as well as animation.

Preparing workspace

Our new dependency for this notebook is Newtonsoft.Json, since json is so easy to use from javascript. I am omitting the bulk of the workspace preparation boilerplate, if you are interested you can check the jupyter notebook directly here.


open CNTK
DeviceDescriptor.UseDefaultDevice().Type
|> printfn "Congratulations, you are using CNTK for: %A" 
Output:
Congratulations, you are using CNTK for: CPU
let device = CNTK.DeviceDescriptor.CPUDevice
let dataType = CNTK.DataType.Float
let initialization = CNTKLib.GlorotUniformInitializer(1.0)

#load "../fsx/CntkHelpers.fsx"

Since we won't be doing any training in this notebook, updating the PATH environment variable is unnecessary.

If however you do happen to try to train a model in this network, this is why it will fail. For more info on the subject check the Preparing workspace notebook (blog) (GitHub)

Building a sample model

Let's create a simple classifier with one hidden layer of 50 nodes that outputs to a softmax function:

open CntkHelpers
let input_dim, num_output_classes = 2,2
let input = Variable.InputVariable(shape [|input_dim|], dataType, "Features")
let label = Variable.InputVariable(shape [|num_output_classes|], dataType, "Labels")
let z = 
    let layers = fullyConnectedClassifierNet input [50] num_output_classes CNTKLib.Sigmoid
    CNTKLib.CrossEntropyWithSoftmax(Var layers, label, "CrossEntropyWithSoftmax")    

Not too complex and not too simple, we will be using this fully connected hidden layer network as our main example when visualizing graphs.

The Search for Nodes

Out of the box, the only way to retrieve the model's constituent CNTK.Function objects seems to be get them by name, using Function.FindByName/FindAllByName. Maintaining a naming scheme while building a model in CNTK is in my opinion not especially expedient at the best of times, and a downright pain if you are looking to implement functional idioms.

In addition, if you are trying to catch a glimpse of the internal workings of a CNTK model it's likely because you had no hand in creating it in the first place, so devising a way to blindly extract network architecture seems worthwhile.

We will do this by implementing a simple graph search algorithm that travels from each Function to its connected Variables, all the while using the Owner property of each Variable to find the rest of the model's component Functions and recursively explore them in turn.

/// Graph search for CNTK.Function objects.
/// We use the Input & Output lists in a Function
/// to find other Function objects by checking
/// the Variable's Owner property.
/// <remarks> CNTK helper function </remarks>
let decomposeFunction (root: Function) = 
    let visited = System.Collections.Generic.Dictionary<string, Function>()            

    let rec expand (f: Function) = 
        match visited.ContainsKey(f.Uid) with
        | true -> Seq.empty
        | false -> 
            visited.Add(f.Uid, f)
            seq {
                yield f
                yield! 
                    seq { yield! f.Inputs
                          yield! f.Outputs }
                    |> Seq.map (fun v -> v.Owner)      
                    |> Seq.filter (not<<isNull)
                    |> Seq.collect expand
            }        

    Array.ofSeq (expand root)
// Example: 
decomposeFunction z |> Array.map (fun f -> f.AsString()) |> Array.rev
Output:
[|"Times: Input('Features', [2], [*, #]) -> Output('@', [50], [*, #])";
  "Plus: Output('@', [50], [*, #]) -> Output('+', [50], [*, #])";
  "StableSigmoid: Output('+', [50], [*, #]) -> Output('StableSigmoid10_Output_0', [50], [*, #])";
  "Times: Output('StableSigmoid10_Output_0', [50], [*, #]) -> Output('@', [2], [*, #])";
  "Plus: Output('@', [2], [*, #]) -> Output('+', [2], [*, #])";
  "CrossEntropyWithSoftmax: Output('+', [2], [*, #]), Input('Labels', [2], [*, #]) -> Output('CrossEntropyWithSoftmax', [1], [*, #])";
  "Composite(CrossEntropyWithSoftmax): Input('Features', [2], [*, #]), Input('Labels', [2], [*, #]) -> Output('CrossEntropyWithSoftmax', [1], [*, #])"|]

This is a brute force algorithm. Perhaps as the complexity of our models increases we will have to revise, but for now it does fine.

Converting to dot notation

The great thing about GraphViz is that it does all the heavy lifting of actually building and arranging the graph, as long as we provide it with pairs of connected nodes (in no particular order). Any additional styling information such as the shape of a node can be declared separately.

Creating dot notation for a CNTK model is as simple as enumerating the component functions and declare edges according to the content of each CNTK.Function object's Input & Output Variable list, with an additional pass to provide styling for each node according to its properties.

For styling I mostly referred to the original CNTK graph logging implementation.

/// Styling functions
module GraphNodeStyle =
    let varText (v:Variable) = if String.IsNullOrEmpty v.Name then v.Uid else v.Name
    let funText (f: Function) = if String.IsNullOrEmpty f.Name then f.Uid else f.Name

    let varLabel (v: Variable) = sprintf "%s [label=\"%s\"];" v.Uid (varText v)
    let funLabel (f: Function) = sprintf "%s [label=\"%s\"];" f.Uid (funText f)
    let edgeLabel (v : Variable) = 
        v.AsString()
         .Replace("(","\n").Replace(")","")
         .Replace("'","").Replace("->","\n->\n")
        |> sprintf "[label=\"%s\"]"            

    let varShape (v: Variable) =
        match v with
        | _ when v.IsInput -> sprintf "%s [shape=invhouse, color=yellow];" v.Uid
        | _ when v.IsOutput -> sprintf "%s [shape=invhouse, color=gray];" v.Uid
        | _ when v.IsPlaceholder -> sprintf "%s [shape=invhouse, color=yellow];" v.Uid
        | _ when v.IsParameter -> sprintf "%s [shape=diamond, color=green];" v.Uid
        | _ when v.IsConstant -> sprintf "%s [shape=rectangle, color=lightblue];" v.Uid
        | _ -> sprintf "%s [shape=circle, color=purple];" v.Uid

    let funShape (f: Function) = 
        match f with 
        | _ when f.IsComposite -> sprintf "%s [shape=ellipse, fontsize=20, penwidth=2, peripheries=2];" f.Uid
        | _ when f.IsPrimitive -> sprintf "%s [shape=ellipse, fontsize=20, penwidth=2, size=0.6];" f.Uid
        | _ -> sprintf "%s [shape=ellipse, fontsize=20, penwidth=4];" f.Uid

    let varEdges (f: Function) (v: Variable) = 
        let inputIndex = f.Inputs |> Seq.map (fun v -> v.Uid) |> Set
        let outputIndex = f.Outputs |> Seq.map (fun v -> v.Uid) |> Set

        match inputIndex.Contains(v.Uid), outputIndex.Contains(v.Uid) with 
        | true, _ when v.IsParameter -> sprintf "%s -> %s %s;" v.Uid f.Uid (edgeLabel v) |> Some
        | _, true when v.IsParameter -> sprintf "%s -> %s [label=\"output param\"];" f.Uid v.Uid|> Some
        | true, _ -> sprintf "%s -> %s %s;" v.Uid f.Uid (edgeLabel v) |> Some 
        | _, true -> sprintf "%s -> %s [label=\"output\"];" f.Uid v.Uid |> Some            
        | _ -> None

    let varOwner (v: Variable) =
        match v.Owner with
        | null -> None
        | _ -> sprintf "%s -> %s [style=\"dotted\"];" v.Owner.Uid v.Uid |> Some
/// Mapping the connections to and from a CNTK.Function object 
/// in GraphViz dot notation.
/// https://github.com/Microsoft/CNTK/blob/master/bindings/python/cntk/logging/graph.py
/// <remarks> CNTK helper function </remarks>
let extractGraphVizDotNotation (f: Function) =     
    let vars = Seq.append f.Inputs f.Outputs
    let funs = seq { 
            yield f
            yield f.RootFunction;
            yield! vars |> Seq.map (fun v -> v.Owner) |> Seq.filter (isNull>>not) 
        } 

    seq {        
        if f.Uid <> f.RootFunction.Uid 
        then yield sprintf "%s -> %s [label=\"root function\"];" f.RootFunction.Uid f.Uid
        yield! vars |> Seq.map GraphNodeStyle.varShape
        yield! vars |> Seq.map GraphNodeStyle.varLabel
        yield! vars |> Seq.map (GraphNodeStyle.varEdges f) |> Seq.choose id
        //yield! vars |> Seq.map varOwner |> Seq.choose id
        yield! funs |> Seq.map GraphNodeStyle.funLabel 
        yield! funs |> Seq.map GraphNodeStyle.funShape
    } |> Seq.distinct

Using GraphViz inside an F# jupyter notebook

As far as I can tell there are restrictions on loading javascript by using the Util.Html command. I was not able to initialize d3-graphviz either by including the cdn reference or by downloading the library files and referring to them locally.

What did work was using webpack to bundle the d3-graphviz library files together, along with javascript event handlers to serve as API endpoints; which is to say, instead of directly calling the d3.graphviz object we trigger an event with the dot notation as payload, which will then be handled from the bundled js.

This seems like a very roundabout way of doing things, and if a better way to integrate tons of html/js with IfCntk comes up I will be glad to revisit and hopefully provide a more elegant solution.

The webpack project can be found here.

@"<script src='../../../d3-jupyter/dist/bundle.js'></script>" 
|> Util.Html |> Display

While I include the bundle.js in the repository, you should remember that making changes to the index.js requires you to rebuild the project with webpack.

Wrapper functions for accessing d3-graphviz from the notebook

In lieu of a normal API.

/// Provide jQuery paths to the html elements
/// you want the model graph and the node 
/// descriptions to appear for any subsequent
/// calls to renderDor and renderSeries
let initGraph infoPath graphPath = 
    sprintf "<script>$(document).trigger('INIT_D3', ['%s','%s']);</script>" infoPath graphPath
    |> Util.Html


/// Send dot notation to be rendered as a graph
let renderDot engine dotNotation jsonInfo = 
    sprintf "<script>$(document).trigger('RENDER_GRAPH', [`%s`,`%s`, '%s']);</script>" dotNotation jsonInfo engine
    |> Util.Html

/// Send a series of graphs to be 
/// sequentially animated into each other
let renderSeries engine (digraphs : string[]) = 
    digraphs
    |> Array.reduce(sprintf "%s','%s")
    |> fun graphs -> graphs,engine
    ||> sprintf "<script>$(document).trigger('RENDER_SERIES', [['%s'],'%s']);</script>" 
    |> Util.Html   

Bringing the whole thing together

/// Extract notation for each CNTK.Function in model
/// combine in a string and send to graphviz for rendering
let displayGraphForModel hostId (model: Function) =
    let dotNotation =
        model
        |> decomposeFunction
        |> Array.filter (fun f -> not f.IsComposite)
        |> Array.collect (extractGraphVizDotNotation>>Array.ofSeq)
        |> Array.distinct 
        |> Array.reduce(sprintf "%s\n%s")
        |> sprintf "digraph { %s }"
            
    initGraph "" ("#"+hostId) |> Display
    renderDot "dot" dotNotation "" |> Display   

In order to be able to show the d3-graphviz output we need a host for the svg inside the notebook. This is trivial to do in IfCntk, again by using the Util.Html:

let displayGraphHost hostId =
    """
<div id="hostId" 
     style="width: 100%; border: solid lightblue 1px; border-radius: 15px; position: relative">    
</div>""".Replace("hostId", hostId)
    |> Util.Html |> Display
displayGraphHost "graph"
displayGraphForModel "graph" z
Output:
graphviz dot layout 2x50x2 Sigmoid/CrossEntropy /w SoftMax

Feel free to exclude or include the composite function node in your graph, since it is seems generally redundant for so simple a model, where it's implied that all the component nodes are parts of a composite model anyway.

It should be noted that CNTK's own Python implementation ommits composite function nodes completely.

Differences from the Python version

The most noticable difference so far seems to be that the result of the CNTK.logging.graph.plot of the Python API ommits the Function -> Variable owner relationship (which in this image is additionally denoted with a dotted edge), and that there seems to be no option to include the composite function in the graph - not that you would want to in the first place; as I noted earlier, it only seems to make the graph unnecessarily convoluted -- on the other hand, in more complex models the ability to show parts of the network as a composite node with the option to expand is bound to come in handy.

Similarly, the reason I am keeping the owner edges for now is less for completeness' sake and more because it easier than rewriting the code to prune them, so hopefully they will also prove useful in the long run and save me the refactoring...

One more thing about the owner edges, unlike most other cases I do not label them with the content of CNTK.Function.AsString() since it is very wordy and mainly just a restatement of the input and output Variables' .AsString() descriptions.

Also, whereas the python version always uses a default property to label each node, here we default to Uid only if no explicit name has been set. Thus, while it may seem that I have maintained the convention of denoting dot product nodes with @ and addition nodes with +, in fact they are just named that way in the fullyConnectedClassifierNet. The same is true for the cross entropy / softmax node.

Let's add some animation

d3-graphviz allows us to animate one graph into an other, simply by declaring a transition filter and making consecutive calls to the render function. Best results require consequent renderings be done after the current graph has finished rendering, so we cannot just call renderGraph a bunch of times and be done with it, which is why the RENDER_SERIES event is used.

For javascript implementation details see index.js

/// Extract dot notation by starting with an empty model
/// and then adding a "frame" of dot notation for
/// each additional CNTK.Function in the model.
let displayGradualGraph hostId (model: Function) =

    initGraph "" ("#" + hostId) |> Display

    let dotNotation =
        model
        |> decomposeFunction 
        //|> Array.filter (fun f -> not f.IsComposite)
        |> Array.map (extractGraphVizDotNotation>>Seq.reduce(sprintf "%s %s")) 
        |> Array.rev
        |> Array.scan (sprintf "%s %s") ""
        |> Array.map (sprintf "digraph { %s }")
        |> Array.map (fun txt -> txt.Replace("\n","\\n"))
           
    renderSeries "dot" dotNotation
    |> Display
displayGraphHost "gradualGraph"
displayGradualGraph "gradualGraph" z
Output:

One other use I have found for including the composite Function node in the graph, is that it does make the final step of the animation sequence more impressive, offering a finale of sorts.

Adding an information panel next to the graph

d3js makes it pretty easy to hang events any element in your graph, so why not present a property dump for each node on mouseover, along with the graph?

First we'll need a function to retrieve said property dump. Most CNTK objects implement an .AsString() function that returns a very short summary of the item, so we will use that when it's available. There also needs to be some special handling for enumerable properties, and for the rest we can just call .ToString().

open System.Collections.Generic

/// Active pattern to separate properties for special handling.
let (|IsEnumerable|IsDescribable|IsPrimitive|) (t: Type) = 
    if  t <> typeof<string> && 
        (typeof< IEnumerable<_> >).IsAssignableFrom(t) 
    then IsEnumerable
    else if t.GetMethods() 
            |> Array.exists (fun meth -> meth.Name = "AsString") 
    then IsDescribable
    else IsPrimitive

/// Helper function to convert any property to string,
/// always checking if .AsString() is available
let asString item =
    if isNull item then ""
    else
        match item.GetType() with
        | IsDescribable ->     
            item.GetType()
                .GetMethod("AsString", Array.empty)
                .Invoke(item, Array.empty).ToString()
        | _ -> item.ToString()
open System.Reflection

/// A simple property serializer for CNTK nodes
/// <remarks> CNTK Helper function </remarks>
let describeNode (item: obj) =
    [|
        yield KeyValuePair("NodeType", item.GetType().Name)
        yield KeyValuePair("AsString", item |> asString)
        yield!
            item.GetType().GetProperties()
            |> Seq.map
                (fun prop ->
                    match prop.PropertyType with
                    | IsEnumerable -> 
                        prop.Name, 
                          (prop.GetValue(item) :?> IEnumerable<_>) 
                          |> function 
                          | list when list |> Seq.isEmpty -> "[]" 
                          | list -> 
                              list 
                              |> Seq.map (asString)                         
                              |> Seq.reduce (sprintf "%s, %s")
                    | IsDescribable -> prop.Name, prop.GetValue(item) |> asString
                    | IsPrimitive -> 
                        prop.Name,
                            try prop.GetValue(item) |> asString
                            with ex -> sprintf "%s" ex.Message)
            |> Seq.map (KeyValuePair)
    |]
Mapping the result of describeNode to KeyValuePair is just there to make the eventual json serialization that we send to the javascript bundle a bit more readable.

In order to display the node info table properly, we need expand our html host accordingly:

let displayGraphHostWithInfoPanel graphHostId infoPanelId =
    """
<div style="width: 100%; height: 700px; max-height: 700px; overflow-y: scroll; overflow-x: hidden; border: solid lightblue 1px; border-radius: 15px; position: relative">
    <div id="graphHostId" style="width: 100%; height: 100%;"></div>
    <div id="infoPanelId" style="width: 350px; max-width: 350px; position: absolute; top: 0; right: 0; background-color: white"></div>
</div>
    """.Replace("graphHostId",graphHostId).Replace("infoPanelId",infoPanelId)
    |> Util.Html |> Display

Then, we expand out dot notation builder with the code that extracts property information for each node. In order to keep things from becoming to convoluted, we extract all the node descriptions in one pass, and send them to the bundled javascript all at once.

open Newtonsoft.Json

let displayGraphWithInfo graphHostId infoHostId (model: Function) =
    let funcs = model |> decomposeFunction |> Array.filter (fun f -> not f.IsComposite)
    let vars = 
        funcs 
        |> Array.collect(fun f -> [|yield! f.Inputs; yield! f.Outputs|])
        |> Array.distinctBy (fun v -> v.Uid)
    
    let describe (tensor: obj) =
        tensor |> describeNode |> JsonConvert.SerializeObject
    
    // This produces a json object where each property is named with
    // the corresponding node's Uid and contains the property dump 
    // produced by the describeNode function.
    let nodeInfo = 
        [|  yield "{"
            yield
                funcs 
                |> Array.map 
                    (fun f -> sprintf "\"%s\": %s" (f.Uid) (describe f))
                |> Array.reduce(sprintf "%s,\n%s")
            yield ","
            yield
                vars 
                |> Array.map 
                    (fun v -> sprintf "\"%s\": %s" (v.Uid) (describe v))
                |> Array.reduce(sprintf "%s,\n%s")
            yield "}" |]
        |> Array.reduce(sprintf "%s\n%s")
        |> fun text -> text.Trim();
    
    let dotNotation =
        funcs
        |> Array.collect (extractGraphVizDotNotation>>Array.ofSeq)
        |> Array.distinct 
        |> Array.reduce(sprintf "%s\n%s")
        |> sprintf "digraph { %s }"
            
    initGraph ("#" + infoHostId) ("#" + graphHostId) |> Display
    renderDot "dot" dotNotation nodeInfo |> Display
displayGraphHostWithInfoPanel "graphWithPanel" "infoPanel"
displayGraphWithInfo "graphWithPanel" "infoPanel" z
Output:
graph with node information panel from notebook

Additional graph layout engines

module Engine =    
    let dot = "dot"
    let circo = "circo"
    let fdp = "fdp"
    let neato = "neato"
    let osage = "osage"
    let patchwork = "patchwork"
    let twopi = "twopi"

Like Graphviz, d3-graphviz allows for a small variety of different layout engines to draw its graphs.

Originally I intended make it possible to exchange engines on the fly via combobox selection (you may have noticed that the wrapper functions for javascript calls still expose an engine parameter), until I realized that CNTK models look mostly terrible with every option but the default of dot...

You're welcome to try them of course. Here are a few samples of what to expect:

dot
(default)
circo
fdp
neato
osage
patchwork
twopi

Thanks for making it to the end! The image recognition notebook is probably coming next, although I've been spending a lot of time with Fabulous lately, so we'll see. Hopefully I'll also enable article comments soon enough, it's entirely a matter of mustering enough energy to go through disqus'  documentation on GDPR, in the meantime feel free to hit me up on twitter. Cheers! -- Ares

 

Links in article