Edge-Computing-Engine/neral/test.py

class Tensor:
    def __init__(self, data, depend=[]):
        """初始化"""
        self.data = data  
        self.depend = depend 
        self.grad = 0 
    def __mul__(self, data):
        """乘法"""
        def grad_fn1(grad):
            return grad * data.data 
        def grad_fn2(grad):
            return grad * self.data  
        depend = [(self, grad_fn1), (data, grad_fn2)]
        new = Tensor(self.data * data.data, depend) 
        return new 
    def __rmul__(self, data):
        def grad_fn1(grad):
            return grad * data.data 
        def grad_fn2(grad):
            return grad * self.data  
        depend = [(self, grad_fn1), (data, grad_fn2)]
        new = Tensor(self.data * data.data, depend) 
        return new 
    def __add__(self, data):
        """加法"""
        def grad_fn(grad):
            return grad   
        depend = [(self, grad_fn), (data, grad_fn)]
        new = Tensor(self.data * data.data, depend) 
        return new    
    def __radd__(self, data):
        def grad_fn(grad):
            return grad   
        depend = [(self, grad_fn), (data, grad_fn)]
        new = Tensor(self.data * data.data, depend) 
        return new  
    def __repr__(self):
        return f"Tensor:{self.data}"
    def backward(self, grad=None):
        """
        反向传播，需要递归计算
        """
        if grad == None:
            self.grad = 1 
        else:
            # 这一步用于计算图中的分支
            self.grad += grad
        # 这一步是递归计算
        for tensor, grad_fn in self.depend:
            bw = grad_fn(self.grad)
            tensor.backward(bw)
x = Tensor(4) 
f = x * x 
g = x * x 
y = f + g
y.backward()
print(x)
print(y, g.grad, x.grad)