sgd_hw/layer.py

224 lines
6.1 KiB
Python
Raw Normal View History

2021-02-12 15:50:20 +09:00
import numpy as np
from contextlib import contextmanager
from typing import Generator, Dict, Union
import io
#only scalar gradient
#op must be tree. 그래프 구현할려면, 위상정렬해서 순회해야하기 때문에 그렇게 하지 않음.
2021-02-13 00:26:53 +09:00
def broadcasting_be(a,b):
i = len(a)-1
j = len(b)-1
abroad = []
bbroad = []
while i >= 0 and j >= 0:
if a[i] == b[j]:
abroad.insert(0,1)
bbroad.insert(0,1)
elif a[i] == 1 or b[j] == 1:
abroad.insert(0,b[j])
bbroad.insert(0,a[i])
else:
raise ValueError
i -= 1
j -= 1
while i >= 0:
bbroad.insert(0,a[i])
i -= 1
while j >= 0:
abroad.insert(0,b[j])
j -= 1
return abroad, bbroad
2021-02-12 15:50:20 +09:00
class NonExistVarableError(ValueError):
pass
def make_mermaid_graph(result):
with io.StringIO("") as graph:
graph.write("graph TD\n")
result.mermaid_graph(graph)
graph.write(f"{id(result)}-->Result\n")
return graph.getvalue()
class OpTree:
def __init__(self):
super().__init__()
def __matmul__(self,a):
return MatMulOp(self,a)
def __add__(self,a):
return AddOp(self,a)
@property
def T(self):
return TransposeOp(self)
class MatMulOp(OpTree):
def __init__(self,a,b):
super().__init__()
self.a = a
self.b = b
va = self.a.numpy() if isinstance(self.a,OpTree) else self.a
vb = self.b.numpy() if isinstance(self.b,OpTree) else self.b
self.v = va @ vb
def __str__(self):
return f"MatmulOp"
def mermaid_graph(self,writer):
if isinstance(self.a,OpTree):
self.a.mermaid_graph(writer)
writer.write(f'{id(self.a)}-->{id(self)}[MatmulOp]\n')
if isinstance(self.b,OpTree):
self.b.mermaid_graph(writer)
writer.write(f'{id(self.b)}-->{id(self)}[MatmulOp]\n')
def numpy(self):
return self.v
def backprop(self,seed):
#a @ b
a = self.a.numpy() if isinstance(self.a,OpTree) else self.a
b = self.b.numpy() if isinstance(self.b,OpTree) else self.b
if isinstance(self.a,OpTree):
s = seed * np.transpose(b) if seed.shape == () else (seed) @ np.transpose(b)
self.a.backprop((s))
if isinstance(self.b,OpTree):
s = np.transpose(a) * seed if seed.shape == () else np.transpose(a) @ seed
self.b.backprop(s)
def matmul(a,b):
return MatMulOp(a,b)
class AddOp(OpTree):
def __init__(self,a,b):
super().__init__()
self.a = a
self.b = b
va = self.a.numpy() if isinstance(self.a,OpTree) else self.a
vb = self.b.numpy() if isinstance(self.b,OpTree) else self.b
self.v = va + vb
def __str__(self):
return f"AddOp"
def mermaid_graph(self,writer):
if isinstance(self.a,OpTree):
self.a.mermaid_graph(writer)
writer.write(f'{id(self.a)}-->{id(self)}[AddOp]\n')
if isinstance(self.b,OpTree):
self.b.mermaid_graph(writer)
writer.write(f'{id(self.b)}-->{id(self)}[AddOp]\n')
def numpy(self):
return self.v
def backprop(self,seed):
#a + b
2021-02-13 00:26:53 +09:00
ashape, bshape = broadcasting_be(self.a.numpy().shape,self.b.numpy().shape)
aai = np.where(np.array(ashape) != 1)
bbi = np.where(np.array(bshape) != 1)
2021-02-12 15:50:20 +09:00
if isinstance(self.a,OpTree):
2021-02-13 00:26:53 +09:00
self.a.backprop(np.sum(seed,axis=tuple(aai[0])))
2021-02-12 15:50:20 +09:00
if isinstance(self.b,OpTree):
2021-02-13 00:26:53 +09:00
self.b.backprop(np.sum(seed,axis=tuple(bbi[0])))
2021-02-12 15:50:20 +09:00
def addmul(a,b):
return AddOp(a,b)
class FunctionOp(OpTree):
def __init__(self,f, f_grad, f_name, i):
super().__init__()
self.f = np.vectorize(f)
self.f_grad = np.vectorize(f_grad)
self.f_name = f_name
self.i = i
self.v = self.f(i.numpy())
def __str__(self):
return f"Function{self.f_name}Op"
def mermaid_graph(self,writer):
self.i.mermaid_graph(writer)
writer.write(f'{id(self.i)}-->{id(self)}[Function{self.f_name}Op]\n')
def numpy(self):
return self.v
def backprop(self,seed):
self.i.backprop(seed * (self.f_grad(self.i.numpy())))
class TransposeOp(OpTree):
def __init__(self, i):
super().__init__()
self.i = i
self.v = np.transpose(i.numpy())
def __str__(self):
return f"TransposeOp"
def mermaid_graph(self,writer):
self.i.mermaid_graph(writer)
writer.write(f'{id(self.i)}-->{id(self)}[TransposeOp]\n')
def numpy(self):
return self.v
def backprop(self,seed):
self.i.backprop(np.transpose(seed))
def transposemul(a):
return TransposeOp(a)
def relu(v):
relu_f = lambda x: np.max([x,0])
relu_diff = lambda x: 1 if x > 0 else 0
return FunctionOp(relu_f,relu_diff,"Relu",v)
2021-02-13 13:20:59 +09:00
#row vector
def softmaxHelp(i):
2021-02-20 15:25:40 +09:00
m = np.max(i,axis=i.ndim-1)
e = np.exp(i-m.reshape(*m.shape,1))
sumofe = np.sum(e,axis=i.ndim - 1)
2021-02-13 13:20:59 +09:00
sumofe = sumofe.reshape(*sumofe.shape,1)
return e / sumofe
class SoftmaxWithNegativeLogLikelihood(OpTree):
#row vector
def __init__(self, i, y):
super().__init__()
2021-02-20 15:25:40 +09:00
epsilon = 1e-12
2021-02-13 13:20:59 +09:00
self.i = i
2021-02-13 23:26:40 +09:00
self.s = softmaxHelp(i.numpy())
self.y = y
2021-02-20 15:25:40 +09:00
self.v = -y*np.log(self.s+epsilon)
2021-02-13 23:26:40 +09:00
self.v = np.sum(self.v,axis=self.v.ndim-1)
2021-02-12 15:50:20 +09:00
2021-02-13 13:20:59 +09:00
def __str__(self):
return f"SoftmaxWithNegativeLogLikelihoodOp"
def mermaid_graph(self,writer):
self.i.mermaid_graph(writer)
writer.write(f'{id(self.i)}-->{id(self)}[SoftmaxWithNegativeLogLikelihoodOp]\n')
def numpy(self):
return self.v
def softmax_numpy(self):
return self.s
def backprop(self,seed):
2021-02-13 23:26:40 +09:00
self.i.backprop(seed * (self.s-self.y))
2021-02-13 13:20:59 +09:00
2021-02-12 15:50:20 +09:00
class Variable(OpTree):
def __init__(self,x):
super().__init__()
self.x = x
self.grad = None
def numpy(self):
return self.x
def mermaid_graph(self,writer):
writer.write(f'{id(self)}["Variable{self.x.shape}"]\n')
def backprop(self,seed):
self.grad = seed