| # Licensed to the Apache Software Foundation (ASF) under one |
| # or more contributor license agreements. See the NOTICE file |
| # distributed with this work for additional information |
| # regarding copyright ownership. The ASF licenses this file |
| # to you under the Apache License, Version 2.0 (the |
| # "License"); you may not use this file except in compliance |
| # with the License. You may obtain a copy of the License at |
| # |
| # http://www.apache.org/licenses/LICENSE-2.0 |
| # |
| # Unless required by applicable law or agreed to in writing, |
| # software distributed under the License is distributed on an |
| # "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY |
| # KIND, either express or implied. See the License for the |
| # specific language governing permissions and limitations |
| # under the License. |
| |
| import JSON |
| |
| """ |
| to_graphviz(network) |
| |
| * `network::SymbolicNode`: the network to visualize. |
| * `title::AbstractString:` keyword argument, default "Network Visualization", |
| the title of the GraphViz graph. |
| * `input_shapes`: keyword argument, default `nothing`. If provided, |
| will run shape inference and plot with the shape information. Should |
| be either a dictionary of name-shape mapping or an array of shapes. |
| |
| Returns the graph description in GraphViz `dot` language. |
| """ |
| function to_graphviz(network :: SymbolicNode; title="Network Visualization", input_shapes=nothing) |
| if !isa(input_shapes, Cvoid) |
| internals = get_internals(network) |
| if isa(input_shapes, Dict) |
| _, out_shapes, _ = infer_shape(internals; input_shapes...) |
| else |
| _, out_shapes, _ = infer_shape(internals, input_shapes...) |
| end |
| @assert(!isa(out_shapes, Cvoid), "Failed to do shape inference, input shapes are incomplete") |
| shape_dict = Dict(zip(list_outputs(internals), out_shapes)) |
| draw_shape = true |
| else |
| draw_shape = false |
| end |
| |
| conf = JSON.parse(to_json(network)) |
| nodes = conf["nodes"] |
| heads = unique([x[1]+1 for x in conf["heads"]]) |
| node_attr = Dict(:shape => :box, :fixedsize => true, :width => 1.3, |
| :height => 0.8034, :style => (:rounded, :filled), :penwidth => 2) |
| io = IOBuffer() |
| println(io, "digraph $(_simple_escape(title)) {") |
| println(io, "node [fontsize=10];") |
| println(io, "edge [fontsize=10];") |
| |
| # color map |
| fillcolors = ("#8dd3c7", "#fb8072", "#ffffb3", "#bebada", "#80b1d3", |
| "#fdb462", "#b3de69", "#fccde5") |
| edgecolors = ("#245b51", "#941305", "#999900", "#3b3564", "#275372", |
| "#975102", "#597d1c", "#90094e") |
| |
| # make nodes |
| for i = 1:length(nodes) |
| node = nodes[i] |
| op = node["op"] |
| name = node["name"] |
| attr = deepcopy(node_attr) |
| label = op |
| |
| # Up to 0.11.0 version of mxnet additional info was stored in |
| # node["attr"]. Staring from 0.12 `attr` was changed to `attrs`. |
| # See: https://github.com/dmlc/nnvm/pull/152 |
| if haskey(node, "attrs") |
| node_info = node["attrs"] |
| elseif haskey(node, "attr") |
| node_info = node["attr"] |
| end |
| |
| if op == "null" |
| if i ∈ heads |
| # heads are output nodes |
| label = node["name"] |
| colorkey = 1 |
| else |
| # otherwise, input nodes, might be data, label or parameters |
| continue |
| end |
| elseif op == "Convolution" |
| if haskey(node_info,"stride") |
| stride_info=_extract_shape(node_info["stride"]) |
| else |
| stride_info="1" |
| end |
| |
| label = format("Convolution\nkernel={1}\nstride={2}\nn-filter={3}", |
| _extract_shape(node_info["kernel"]), |
| stride_info, |
| node_info["num_filter"]) |
| colorkey = 2 |
| elseif op == "FullyConnected" |
| label = format("FullyConnected\nnum-hidden={1}", node_info["num_hidden"]) |
| colorkey = 2 |
| elseif op == "Activation" |
| label = format("Activation\nact-type={1}", node_info["act_type"]) |
| colorkey = 3 |
| elseif op == "BatchNorm" |
| colorkey = 4 |
| elseif op == "Pooling" |
| if haskey(node_info,"stride") |
| stride_info=_extract_shape(node_info["stride"]) |
| else |
| stride_info="1" |
| end |
| label = format("Pooling\ntype={1}\nkernel={2}\nstride={3}", |
| node_info["pool_type"], |
| _extract_shape(node_info["kernel"]), |
| stride_info) |
| colorkey = 5 |
| elseif op ∈ ("Concat", "Flatten", "Reshape") |
| colorkey = 6 |
| elseif endswith(op, "Output") || op == "BlockGrad" |
| colorkey = 7 |
| else |
| colorkey = 8 |
| end |
| |
| if op != "null" |
| label = "$name\n$label" |
| end |
| attr[:fillcolor] = fillcolors[colorkey] |
| attr[:color] = edgecolors[colorkey] |
| attr[:label] = label |
| _format_graphviz_node(io, name, attr) |
| end |
| |
| # add edges |
| for i = 1:length(nodes) |
| node = nodes[i] |
| op = node["op"] |
| name = node["name"] |
| if op == "null" |
| continue |
| end |
| inputs = node["inputs"] |
| for item in inputs |
| input_node = nodes[item[1]+1] |
| input_name = input_node["name"] |
| if input_node["op"] != "null" || (item[1]+1) ∈ heads |
| attr = Dict(:dir => :back, :arrowtail => :open, :color => "#737373") |
| if draw_shape |
| if input_node["op"] != "null" |
| key = Symbol(input_name, "_output") |
| shape = shape_dict[key][1:end-1] |
| else |
| key = Symbol(input_name) |
| shape = shape_dict[key][1:end-1] |
| end |
| label = "(" * join([string(x) for x in shape], ",") * ")" |
| attr[:label] = label |
| end |
| _format_graphviz_edge(io, name, input_name, attr) |
| end |
| end |
| end |
| println(io, "}") |
| |
| return String(take!(io)) |
| end |
| |
| function _format_graphviz_attr(io::IOBuffer, attrs) |
| label = get(attrs, :label, nothing) |
| if isa(label, Cvoid) |
| print(io, " [") |
| else |
| print(io, " [label=$(_simple_escape(label)),") |
| end |
| first_attr = true |
| for (k,v) in attrs |
| if k != :label |
| if !first_attr |
| print(io, ",") |
| end |
| first_attr = false |
| |
| if isa(v, AbstractString) && v[1] == '#' |
| # color |
| v = _simple_escape(v) |
| elseif isa(v, Tuple) |
| v = _simple_escape(join([string(x) for x in v], ",")) |
| end |
| print(io, "$k=$v") |
| end |
| end |
| println(io, "];") |
| end |
| function _simple_escape(str) |
| str = replace(string(str), r"\n" => "\\n") |
| return "\"$str\"" |
| end |
| function _format_graphviz_node(io::IOBuffer, name::AbstractString, attrs) |
| print(io, "$(_simple_escape(name)) ") |
| _format_graphviz_attr(io, attrs) |
| end |
| function _format_graphviz_edge(io::IOBuffer, head, tail, attrs) |
| print(io, """$(_simple_escape(head)) -> $(_simple_escape(tail)) """) |
| _format_graphviz_attr(io, attrs) |
| end |
| function _extract_shape(str :: AbstractString) |
| shape = matchall(r"\d+", str) |
| shape = reverse(shape) # JSON in libmxnet has reversed shape (column vs row majoring) |
| return "(" * join(shape, ",") * ")" |
| end |
