107 lines
3.3 KiB
Python
107 lines
3.3 KiB
Python
import numpy as np
|
|
from activations import ActivationFunc, Identity
|
|
|
|
|
|
class NNLayer:
|
|
def __init__(self,
|
|
in_size: int,
|
|
out_size: int,
|
|
lr: float,
|
|
activation_func: ActivationFunc):
|
|
limit = np.sqrt(6 / (in_size + out_size))
|
|
self.W = np.random.uniform(-limit, limit, (in_size, out_size))
|
|
self.B = np.zeros((out_size))
|
|
self.lr = lr
|
|
self.input = None
|
|
self.output = None
|
|
self.output_linear = None
|
|
self.activation_func = activation_func
|
|
|
|
def __str__(self):
|
|
return f'[ {self.W.shape[0]} => {self.W.shape[1]}\tlr:{self.lr}\tactivation:{self.activation_func.__class__.__name__} ]' # noqa
|
|
|
|
def forward(self, v: np.ndarray) -> np.ndarray:
|
|
self.input = v
|
|
self.output_linear = self.input @ self.W + self.B
|
|
self.output = self.activation_func(
|
|
self.output_linear
|
|
)
|
|
return self.output
|
|
|
|
def backprop(self, error: np.ndarray) -> np.ndarray:
|
|
error *= self.activation_func.d(self.output_linear)
|
|
ret = self.W @ error
|
|
dW = np.outer(self.input, error) * self.lr
|
|
dB = error * self.lr
|
|
self.W -= dW
|
|
self.B -= dB
|
|
return ret
|
|
|
|
|
|
class SampleLayer:
|
|
def __init__(self,
|
|
in_size: int,
|
|
lr: float,
|
|
activation_func: ActivationFunc):
|
|
self.input = None
|
|
self.mean_nn = NNLayer(
|
|
in_size,
|
|
in_size,
|
|
lr,
|
|
activation_func)
|
|
self.std_nn = NNLayer(
|
|
in_size,
|
|
in_size,
|
|
lr,
|
|
activation_func)
|
|
|
|
def DKL(self):
|
|
return -0.5 * np.mean(1 + self.logvar - self.mean ** 2 - np.exp(self.logvar)) # noqa
|
|
|
|
def forward(self, v: np.ndarray) -> np.ndarray:
|
|
self.input = v
|
|
self.mean = self.mean_nn.forward(v)
|
|
self.logvar = np.clip(self.std_nn.forward(v), -10, 10)
|
|
self.std = np.exp(0.5 * self.logvar)
|
|
self.eps = np.random.normal(0, 1, self.mean.shape)
|
|
return 0.5 * self.eps * self.std + self.mean
|
|
|
|
def backprop(self, error: np.ndarray) -> np.ndarray:
|
|
dmean = error + self.mean
|
|
dstd = error * self.eps + 0.5 * (np.exp(self.logvar) - 1)
|
|
mean_error = self.mean_nn.backprop(dmean)
|
|
logvar_error = self.std_nn.backprop(dstd * self.std)
|
|
return mean_error + logvar_error
|
|
|
|
|
|
class DeepNNLayer:
|
|
def __init__(self,
|
|
layers: list[int],
|
|
lr: float,
|
|
activation_func: ActivationFunc):
|
|
self.layers: list[NNLayer] = []
|
|
for i in range(len(layers) - 1):
|
|
self.layers.append(
|
|
NNLayer(
|
|
layers[i],
|
|
layers[i+1],
|
|
lr,
|
|
activation_func if i != len(layers) - 2 else Identity()
|
|
)
|
|
)
|
|
self.in_size = layers[0]
|
|
self.out_size = layers[-1]
|
|
|
|
def __str__(self):
|
|
return '\n'.join([str(layer) for layer in self.layers])
|
|
|
|
def forward(self, v: np.ndarray) -> np.ndarray:
|
|
for layer in self.layers:
|
|
v = layer.forward(v)
|
|
return v
|
|
|
|
def backprop(self, error: np.ndarray) -> np.ndarray:
|
|
for layer in self.layers[::-1]:
|
|
error = layer.backprop(error)
|
|
return error
|