Skip to content
Permalink
Browse files

Move plotting dependencies to test time (#128)

  • Loading branch information
parasj committed Mar 23, 2020
1 parent 3f7caa0 commit b9f1d0c84e3bf8c71394c65c4e399e616c9d0864
Showing with 1 addition and 33 deletions.
  1. +0 −30 checkmate/plot/graph_plotting.py
  2. +1 −3 setup.py
@@ -9,36 +9,6 @@
from checkmate.core.utils.definitions import PathLike


# TODO (paras) fix this function
def tensor_plot(g: DFGraph, sched: Schedule, directory, tag=None, format="pdf", quiet=True):
dot = Digraph("!TensorPlot_{}".format(tag), engine="dot")
if sched is None:
return
for op in sched:
if isinstance(op, OperatorEvaluation):
if g.is_forward_node(op.id):
node_name = g.node_names.get(op.id)
node_name = node_name if node_name is None else "{} ({})".format(node_name, str(op.id))
elif g.is_backward_node(op.id):
node_name = "Grad {}".format(op.id)
else:
raise ValueError("Unknown operation")
# dot.node("op{}".format(op.id), node_name, shape="diamond")
# dot.edge("op{}".format(op.id), "reg{}".format(op.out_register))
dot.node("reg{}".format(op.out_register), "Register {} for {}".format(op.out_register, node_name), shape="box")
for dep_op, dep_reg in op.arg_regs.items():
dot.edge(
"reg{}".format(dep_reg),
"reg{}".format(op.out_register),
style="dashed",
label=str(g.args[op.id].index(dep_op)),
)
try:
dot.render(directory=directory, format=format, quiet=quiet)
except TypeError:
dot.render(directory=directory, format=format)


def plot_dfgraph(g: DFGraph, directory, format="pdf", quiet=True, name=""):
"""Generate Graphviz-formatted edge list for visualization, and write pdf"""
dot = Digraph("render_dfgraph" + str(name))
@@ -8,12 +8,10 @@
python_requires=">=3.5",
install_requires=[
# "keras_segmentation @ https://github.com/ajayjain/image-segmentation-keras/archive/master.zip#egg=keras_segmentation-0.2.0remat",
"matplotlib", # this is only used once in the core checkmate package
"graphviz",
"numpy",
"pandas",
"toposort",
"psutil",
],
extras_require={"test": ["pytest", "tensorflow>=2.0.0"]},
extras_require={"test": ["pytest", "tensorflow>=2.0.0", "matplotlib", "graphviz"]},
)

0 comments on commit b9f1d0c

Please sign in to comment.
You can’t perform that action at this time.