add attributes to nodes

src/ngraph/visualize.cpp

@@ -34,11 +34,43 @@ void Visualize::add(node_ptr p)
     traverse_nodes(p, [&](node_ptr node) {
         for (auto arg : node->get_arguments())
         {
-            m_ss << "    " << arg->get_node_id() << " -> " << node->get_node_id() << ";\n";
+            m_ss << add_attributes(arg);
+            m_ss << add_attributes(node);
+            m_ss << "    " << arg->get_node_id() << " -> " << node->get_node_id();
+            m_ss << ";\n";
         }
     });
 }
 
+std::string Visualize::add_attributes(node_ptr node)
+{
+    string rc;
+    if (!contains(m_nodes_with_attributes, node))
+    {
+        m_nodes_with_attributes.insert(node);
+        rc = get_attributes(node);
+    }
+    return rc;
+}
+
+std::string Visualize::get_attributes(node_ptr node)
+{
+    stringstream ss;
+    if (node->is_parameter())
+    {
+        ss << "    " << node->get_node_id() << " [shape=box color=blue]\n";
+    }
+    else if (node->is_op())
+    {
+        ss << "    " << node->get_node_id() << " [shape=ellipse color=black]\n";
+    }
+    else
+    {
+        ss << "    " << node->get_node_id() << " [shape=diamond color=red]\n";
+    }
+    return ss.str();
+}
+
 void Visualize::save_dot(const string& path) const
 {
 #if GRAPHVIZ_FOUND

src/ngraph/visualize.hpp

@@ -36,6 +36,10 @@ public:
     void save_dot(const std::string& path) const;
 
 private:
-    std::stringstream m_ss;
-    std::string       m_name;
+    std::string add_attributes(node_ptr node);
+    std::string get_attributes(node_ptr node);
+
+    std::stringstream  m_ss;
+    std::string        m_name;
+    std::set<node_ptr> m_nodes_with_attributes;
 };