deepxpy:初步实现构建计算图

Author: miaobyte

front/py/deepx/autograd/__init__.py

@@ -1,6 +1,11 @@
 from .graph import Graph
-from ..tensor.op import Op
+from .node import Node
+from .nodetype import NodeType
+from ._datanode import DataNode
+
 __all__ = [
     'Graph',
-    'Op',
+    'Node',
+    'NodeType',
+    'DataNode',
    ]

front/py/deepx/autograd/_constargnode.py

@@ -0,0 +1,8 @@
+from .node import Node
+from .nodetype import NodeType
+
+
+class ControlFlowNode(Node):
+    def __init__(self, name=None):
+        super().__init__(name="control_flow", ntype=NodeType.CONTROL_FLOW)
+

front/py/deepx/autograd/_controlflownode.py

@@ -1,15 +1,15 @@
-from .node import Node, NodeType
-
+from .node import Node
+from .nodetype import NodeType 
+ 
 class DataNode(Node):
-    def __init__(self, name=None):
-        super().__init__(name=name, ntype=NodeType.TENSOR)
-        self._data = None
+    def __init__(self, name=None, data=None):
+        super().__init__(name=name, ntype=NodeType.DATA)
+        self._data = data
 
     def data(self):
         return self._data
 
     def set_data(self, data):
-        from deepx import Tensor
-        if not isinstance(data, Tensor):
-            raise TypeError("data must be an instance of Tensor")
-        self._data = data
+        self._data = data
+ 
+ 

front/py/deepx/autograd/_datanode.py

@@ -1,133 +1,30 @@
-from .node import Node, NodeType
-
-class OpType:
-    def __init__(self, name, shortchar):
-        self.name = name
-        self.shortchar = shortchar
-    def shortchar(self):
-        return self.shortchar
-# 全局操作类型注册表
-_op_types = {}
-
-def regist_op_type(name, shortchar):
-    """注册一个操作类型"""
-    _op_types[name] = OpType(name, shortchar)
-
-class OpNode(Node):
-    def __init__(self, op_type_name, name=None):
-        if op_type_name not in _op_types:
-            raise ValueError(f"Unknown op type: {op_type_name}")
+from .node import Node
+from .nodetype import NodeType
+
+
+class OpNodeMeta(type):
+    """操作节点元类，负责校验操作名称"""
+    _registered_ops = set()  # 已注册操作名称缓存
+
+    def __call__(cls, name: str, *args, **kwargs):
+        # 在实例化时进行名称校验
+        if name not in cls._registered_ops:
+            raise ValueError(
+                f"Op '{name}' 未注册,请先使用OpNode.register('{name}')注册"
+            )
+        return super().__call__(name, *args, **kwargs)
+
+    @classmethod
+    def register_op(cls, name: str) -> None:
+        """注册新操作类型"""
+        if name in cls._registered_ops:
+            raise ValueError(f"Op '{name}' 已存在")
+        cls._registered_ops.add(name)
+ 
+class OpNode(Node, metaclass=OpNodeMeta):
+    def __init__(self, name: str):
         super().__init__(name=name, ntype=NodeType.OP)
-        self.op_type = _op_types[op_type_name]
-    
-    def shortchar(self):
-        return self.op_type.shortchar
-
-
-regist_op_type("ReLU", "relu")
-regist_op_type("Placeholder", "ph")
-regist_op_type("Neg", "-")
-regist_op_type("Less", "<")
-regist_op_type("Equal", "==")
-regist_op_type("Sigmoid", "σ")
-regist_op_type("Tanh", "tanh")
-regist_op_type("Reshape", "reshape")
-regist_op_type("Transpose", "T")
-regist_op_type("Sum", "Σ")
-regist_op_type("Mean", "μ")
-
-# 操作节点创建函数
-def matmul(a, b, name=None):
-    node = OpNode("MatMul", name)
-    node.add_input("a", a)
-    node.add_input("b", b)
-    return node
-
-def add(a, b, name=None):
-    node = OpNode("Add", name)
-    node.add_input("a", a)
-    node.add_input("b", b)
-    return node
-
-def relu(x, name=None):
-    node = OpNode("ReLU", name)
-    node.add_input("x", x)
-    return node
-
-def placeholder(name=None, shape=None):
-    node = OpNode("Placeholder", name)
-    if shape:
-        node.set_attr("shape", shape)
-    return node
-
-def neg(x):
-    node = OpNode("Neg")
-    node.add_input("x", x)
-    return node
-
-def mul(a, b):
-    node = OpNode("Mul")
-    node.add_input("a", a)
-    node.add_input("b", b)
-    return node
-
-def div(a, b):
-    node = OpNode("Div")
-    node.add_input("a", a)
-    node.add_input("b", b)
-    return node
-
-def sub(a, b):
-    node = OpNode("Sub")
-    node.add_input("a", a)
-    node.add_input("b", b)
-    return node
-
-def less(a, b):
-    node = OpNode("Less")
-    node.add_input("a", a)
-    node.add_input("b", b)
-    return node
-
-def equal(a, b):
-    node = OpNode("Equal")
-    node.add_input("a", a)
-    node.add_input("b", b)
-    return node
-
-def sigmoid(x):
-    node = OpNode("Sigmoid")
-    node.add_input("x", x)
-    return node
-
-def tanh(x):
-    node = OpNode("Tanh")
-    node.add_input("x", x)
-    return node
-
-def reshape(x, shape):
-    node = OpNode("Reshape")
-    node.add_input("x", x)
-    node.set_attr("shape", shape)
-    return node
-
-def transpose(x, dim0, dim1):
-    node = OpNode("Transpose")
-    node.add_input("x", x)
-    node.set_attr("dim0", dim0)
-    node.set_attr("dim1", dim1)
-    return node
-
-def sum(x, dim=None, keepdim=False):
-    node = OpNode("Sum")
-    node.add_input("x", x)
-    node.set_attr("dim", dim)
-    node.set_attr("keepdim", keepdim)
-    return node
 
-def mean(x, dim=None, keepdim=False):
-    node = OpNode("Mean")
-    node.add_input("x", x)
-    node.set_attr("dim", dim)
-    node.set_attr("keepdim", keepdim)
-    return node 
+    @classmethod
+    def register(cls, name: str) -> None:
+        cls.__class__.register_op(name)

front/py/deepx/autograd/_opnode.py

@@ -1,6 +1,8 @@
-from ._datanode import TensorNode
+from ._datanode import DataNode
 from ._opnode import OpNode
-from ._constargnode import ConstArgNode
+from ._controlflownode import ControlFlowNode
+
+eager_mode=False
 
 class Graph:
     # 类属性存储默认实例
@@ -20,33 +22,38 @@ def set_default(cls, graph):
             raise TypeError("Must be a Graph instance")
         cls._default_graph = graph
 
-    def __init__(self):
+    def __init__(self,eager=False):
         self.nodes = []
         self.inputs = []
-        self.tensor_counter = 0  # 添加计数器
-        self.constarg_counter = 0     
-    def add_tensor(self, name, dtype, shape, requires_grad):
-        self.tensor_counter += 1
+        self.data_counter = 0
+        self.control_flow_counter = 0
+        self.eager=eager or eager_mode
+
+    def add_data(self, name, data,inputs=[]):
+        self.data_counter += 1
         if name == "":
-            name = f"tensor_{self.tensor_counter}"
-        node=TensorNode(name, dtype, shape, requires_grad)
-        for input in node.inputs:
-            node.add_input(input.name, input)
+            name = f"data_{self.data_counter}"
+        node=DataNode(name, data)
+        for input in inputs:
+            node.add_input(input)
         self.nodes.append(node)
         return node
-    def add_op(self,name,inputs):
+    def add_op(self,name,inputs=[]):
         node=OpNode(name)
         for input in inputs:
-            node.add_input(input.name, input)
+            node.add_input(input)
         self.nodes.append(node)
+        if self.eager:
+            return node.outputs[0]
         return node
-    def add_constarg(self, value):
-        self.constarg_counter += 1
+    def add_control_flow(self,name,inputs=[]):
+        self.control_flow_counter += 1
         if name == "":
-            name = f"constarg_{self.constarg_counter}"
-        node=ConstArgNode(value)
+            name = f"control_flow_{self.control_flow_counter}"
+        node=ControlFlowNode(name)
+        for input in inputs:
+            node.add_input(input)
         self.nodes.append(node)
         return node
-
 # 初始化默认图
 Graph._default_graph = Graph()

front/py/deepx/autograd/graph.py

@@ -2,7 +2,7 @@
 from .graph import Graph
 from ._datanode import TensorNode
 from ._opnode import OpNode
-from ._constargnode import ConstArgNode
+from ._controlflownode import ConstArgNode
 
 def graph_method(f):
     """装饰器:将函数注册为Graph类的方法"""

front/py/deepx/autograd/graph_viz.py

@@ -1,13 +1,9 @@
 from enum import IntEnum
-
-class NodeType(IntEnum):
-    TENSOR = 0
-    OP = 1 
-    CONST_ARG = 2
-
+from .nodetype import NodeType
+ 
 class Node:
     def __init__(self, 
-                 ntype:NodeType=NodeType.TENSOR,
+                 ntype:NodeType=None,
                    name:str=None,
                    graph =None):
         from .graph import Graph
@@ -18,37 +14,19 @@ def __init__(self,
 
         self._ntype = ntype
         self._name = name
-        self._inputs = {}
-        self.attrs = {}
-    
+        self._inputs = []
+
     @property
     def ntype(self):
         return self._ntype
 
+    @property
+    def graph(self):
+        return self._graph
+    
     @property
     def name(self):
         return self._name
 
-    @property
-    def input(self, name=None):
-        if name is None:
-            return self._inputs
-        else:
-            return self._inputs.get(name)
-    
-    def add_input(self, name, input_node):
-        self._inputs[name] = input_node
-        input_node.outputs.append(self)
-    
-    def remove_input(self, name):
-        if name in self._inputs:
-            node = self._inputs[name]
-            if self in node.outputs:
-                node.outputs.remove(self)
-            del self._inputs[name]
-    
-    def set_attr(self, key, value):
-        self.attrs[key] = value
-        
-    def get_attr(self, key):
-        return self.attrs.get(key) 
+    def add_input(self, input_node):
+        self._inputs.append(input_node)

front/py/deepx/autograd/node.py

@@ -0,0 +1,9 @@
+from enum import IntEnum, EnumMeta
+from typing import Dict, Any
+
+ 
+class NodeType(IntEnum ):
+    DATA = 0
+    OP = 1
+    CONTROL_FLOW = 2
+