{
"cells": [
{
"cell_type": "markdown",
"metadata": {},
"source": [
"# fastdot.pagegraph\n",
"\n",
"> David's stuff!"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"# default_exp pagegraph\n",
"# default_cls_lvl 3"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"#export\n",
"from functools import singledispatch\n",
"from torch import nn"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"#hide\n",
"from nbdev.showdoc import *"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"## Drawing graphs"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"#export\n",
"def split(path, sep = '/'):\n",
" i = path.rfind(sep) + 1\n",
" return path[:i].rstrip(sep), path[i:]"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"#export\n",
"def make_dot_graph(nodes, edges, direction='LR', **kwargs):\n",
" from pydot import Dot, Cluster, Node, Edge\n",
" class Subgraphs(dict):\n",
" def __missing__(self, path):\n",
" parent, label = split(path)\n",
" subgraph = Cluster(path, label=label, style='rounded, filled', fillcolor='#77777744')\n",
" self[parent].add_subgraph(subgraph)\n",
" return subgraph\n",
" g = Dot(rankdir=direction, directed=True, **kwargs)\n",
" g.set_node_defaults(\n",
" shape='box', style='rounded, filled', fillcolor='#ffffff')\n",
" subgraphs = Subgraphs({'': g})\n",
" for path, attr in nodes:\n",
" parent, label = split(path)\n",
" subgraphs[parent].add_node(\n",
" Node(name=path, label=label, **attr))\n",
" for src, dst, attr in edges:\n",
" g.add_edge(Edge(src, dst, **attr))\n",
" return g"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"#export\n",
"def to_dict(inputs):\n",
" return dict(enumerate(inputs)) if isinstance(inputs, list) else inputs"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"#export\n",
"class DotGraph():\n",
" def __init__(self, graph, size=None, direction='LR'):\n",
" self.nodes = [(k, v) for k, (v,_) in graph.items()]\n",
" self.edges = [(src, dst, {'tooltip': name}) for dst, (_, inputs) in graph.items() for name, src in to_dict(inputs).items()]\n",
" self.size, self.direction = size or 8+len(graph)/3, direction\n",
"\n",
" def dot_graph(self, **kwargs):\n",
" return make_dot_graph(self.nodes, self.edges, size=self.size, direction=self.direction, **kwargs)\n",
"\n",
" def svg(self, **kwargs):\n",
" return self.dot_graph(**kwargs).create(format='svg').decode('utf-8')\n",
" try:\n",
" import pydot\n",
" _repr_svg_ = svg\n",
" except ImportError:\n",
" def __repr__(self): return 'pydot is needed for network visualisation'"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"Let's define a compact way of representing directed graphs (suitable for describing neural networks) as python dictionaries. Each node has a *name* which should be a string, a *value* which can be anything and an (ordered) list of *inputs* representing incoming edges from other nodes. Inputs are ordered as they typically represent function args. We can alternatively provide a dict of *inputs* to represent named args. Here is our graph format:\n",
"\n",
"\n",
"```\n",
"Graph := {name: (value, [input])}\n",
"```\n",
"\n",
"or\n",
"\n",
"```\n",
"Graph := {name: (value, {input_name: input}}\n",
"```\n",
"\n",
"Here is an example of such a Graph using both kinds of inputs:\n"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"#define some node values - just strings for now we'll use more interesting values later\n",
"A,B,C,D = 'Aa Bb Cc Dd'.split()\n",
"\n",
"graph = {\n",
" 'a': (A, []),\n",
" 'b': (B, ['a']),\n",
" 'c': (C, {'x': 'a', 'y': 'c'}),\n",
" 'd': (D, ['b'])\n",
"}"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"It will be useful to be able to display Graphs. Here is a function that does that:"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"#export\n",
"def draw_graph(graph):\n",
" return DotGraph({name: ({'tooltip': str(value)}, inputs) for name, (value,inputs) in graph.items()}, size=10)"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [
{
"data": {
"image/svg+xml": [
"\n",
"\n",
"\n",
"\n",
"\n"
],
"text/plain": [
"<__main__.DotGraph at 0x7f31b0d864d0>"
]
},
"execution_count": null,
"metadata": {},
"output_type": "execute_result"
}
],
"source": [
"draw_graph(graph)"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"## Exploding graphs"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"#export\n",
"def iter_nodes(graph):\n",
" # graph = {name: node for (name, node) in graph.items() if node is not None}\n",
" keys = list(graph.keys())\n",
" if not all(isinstance(k, str) for k in keys):\n",
" raise Exception('Node names must be strings.')\n",
" return ((name, (node if isinstance(node, tuple) else (node, [0 if j is 0 else keys[j-1]])))\n",
" for (j, (name, node)) in enumerate(graph.items()))"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"#export\n",
"map_ = lambda func, vals: [func(x) for x in vals] if isinstance(vals, list) else {k: func(v) for k,v in vals.items()}\n",
"pfx = lambda prefix, name: f'{prefix}/{name}'\n",
"external_inputs = lambda graph: set(i for name, (value, inputs) in iter_nodes(graph) for i in inputs if i not in graph)"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"#export\n",
"def bindings(graph, inputs):\n",
" if isinstance(inputs, list): inputs = dict(enumerate(inputs))\n",
" required_inputs = external_inputs(graph)\n",
" missing = [k for k in required_inputs if k not in inputs]\n",
" if len(missing): \n",
" raise Exception(f'Required inputs {missing} are missing from inputs {inputs} for graph {graph}')\n",
" return inputs"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"#export\n",
"walk = lambda dct, key: walk(dct, dct[key]) if key in dct else key"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"#export\n",
"_all_ = ['to_graph']"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"#export\n",
"@singledispatch\n",
"def to_graph(value): \n",
" raise NotImplementedError(f'type = {type(value)}')"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"#export\n",
"@to_graph.register(dict)\n",
"def _(x): return x"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"#export\n",
"def maybe_graph(x):\n",
" try: return to_graph(x)\n",
" except NotImplementedError: return x"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"#export\n",
"def explode(graph, max_levels=-1, convert=maybe_graph):\n",
" graph = convert(graph)\n",
" if max_levels==0 or not isinstance(graph, dict): return graph\n",
" redirects = {}\n",
" def iter_(graph):\n",
" for name, (value, inputs) in iter_nodes(graph):\n",
" value = explode(value, max_levels-1, convert=convert)\n",
" if isinstance(value, dict):\n",
" #special case empty dict\n",
" if not len(value): \n",
" if len(inputs) != 1: raise Exception('Empty graphs (pass-thrus) should have exactly one input')\n",
" redirects[name] = inputs[0] #redirect to input\n",
" else:\n",
" bindings_dict = bindings(value, inputs)\n",
" for n, (val, ins) in iter_nodes(value):\n",
" yield (pfx(name, n), (val, map_((lambda i: bindings_dict.get(i, pfx(name, i))), ins)))\n",
" redirects[name] = pfx(name, n) #redirect to previous node\n",
" else:\n",
" yield (name, (value, inputs))\n",
" return {name: (value, map_((lambda i: walk(redirects, i)), inputs)) for name, (value, inputs) in iter_(graph)}"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"## Creating Networks"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"#export\n",
"class ColorMap(dict):\n",
" palette = ['#'+x for x in (\n",
" 'bebada,ffffb3,fb8072,8dd3c7,80b1d3,fdb462,b3de69,fccde5,bc80bd,ccebc5,ffed6f,1f78b4,33a02c,e31a1c,ff7f00,'\n",
" '4dddf8,e66493,b07b87,4e90e3,dea05e,d0c281,f0e189,e9e8b1,e0eb71,bbd2a4,6ed641,57eb9c,3ca4d4,92d5e7,b15928'\n",
" ).split(',')]\n",
"\n",
" def __missing__(self, key):\n",
" self[key] = self.palette[len(self) % len(self.palette)]\n",
" return self[key]\n",
"\n",
" def _repr_html_(self):\n",
" css = (\n",
" '.pill {'\n",
" 'margin:2px; border-width:1px; border-radius:9px; border-style:solid;'\n",
" 'display:inline-block; width:100px; height:15px; line-height:15px;'\n",
" '}'\n",
" '.pill_text {'\n",
" 'width:90%; margin:auto; font-size:9px; text-align:center; overflow:hidden;'\n",
" '}'\n",
" )\n",
" s = ''\n",
" return ''+''.join((s.format(color, text) for text, color in self.items()))"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"#export\n",
"class Network(nn.Module):\n",
" colors = ColorMap()\n",
" def __init__(self, graph, cache_activations=False):\n",
" self._graph = to_graph(graph)\n",
" super().__init__()\n",
" self.cache_activations = cache_activations\n",
" for path, (val, _) in iter_nodes(self._graph): setattr(self, path.replace('/', '__'), val)\n",
" \n",
" def __setattr__(self, name, value):\n",
" super().__setattr__(name, value)\n",
" path = name.replace('__', '/')\n",
" if path in self._graph:\n",
" old_val = self._graph[path]\n",
" self._graph[path] = (value, old_val[1]) if isinstance(old_val, tuple) else value\n",
"\n",
" def forward(self, *args):\n",
" prev = args[0]\n",
" outputs = self.cache = dict(enumerate(args))\n",
" for k, (node, inputs) in iter_nodes(self._graph):\n",
" if k not in outputs:prev = outputs[k] = node(*[outputs[x] for x in inputs])\n",
" if not self.cache_activations: self.cache = None\n",
" return prev\n",
"\n",
" def draw(self, **kwargs):\n",
" return DotGraph({p: ({'fillcolor': self.colors[type(v).__name__], 'tooltip': str(v)}, inputs)\n",
" for p, (v, inputs) in iter_nodes(to_graph(self))}, **kwargs)\n",
"\n",
" def explode(self, max_levels=-1):\n",
" convert = lambda x: to_graph(x) if isinstance(x, Network) else x\n",
" return Network(explode(self, max_levels, convert=convert))"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"#export\n",
"def to_network(module, max_levels=-1):\n",
" net = Network(module)\n",
" if max_levels == 0: return net\n",
" for k, mod in net.named_children():\n",
" try: setattr(net, k, to_network(mod, max_levels-1))\n",
" except NotImplementedError: pass\n",
" return net"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"#export\n",
"@to_graph.register(Network)\n",
"def _(x): return x._graph "
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"## Export -"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"Converted 00_core.ipynb.\n",
"Converted index.ipynb.\n"
]
}
],
"source": [
"from nbdev.export import notebook2script\n",
"notebook2script()"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": []
}
],
"metadata": {
"kernelspec": {
"display_name": "Python 3",
"language": "python",
"name": "python3"
}
},
"nbformat": 4,
"nbformat_minor": 2
}