pytensorforge 0.1.0__py3-none-any.whl

src/models/transformers/MultiHeadAttention.py

@@ -0,0 +1,130 @@
+from src.core.Tensor import Tensor
+from src.neural.Layer import Layer
+
+
+class MultiHeadAttention(Layer):
+
+    def __init__(
+        self,
+        d_model,
+        num_heads,
+        bias=True,
+    ):
+        super().__init__()
+
+        self.Wv = None
+        self.Wk = None
+        self.Wq = None
+        self.Wo = None
+        if d_model % num_heads != 0:
+            raise ValueError(
+                "d_model must be divisible by num_heads"
+            )
+
+        self.d_model = d_model
+        self.num_heads = num_heads
+        self.head_dim = d_model // num_heads
+
+        self.bias = bias
+
+    def build(self, input_shape):
+
+        self.Wq = self.add_weight(
+            shape=(self.d_model, self.d_model),
+            initializer="glorot_uniform",
+            name="Wq",
+        )
+
+        self.Wk = self.add_weight(
+            shape=(self.d_model, self.d_model),
+            initializer="glorot_uniform",
+            name="Wk",
+        )
+
+        self.Wv = self.add_weight(
+            shape=(self.d_model, self.d_model),
+            initializer="glorot_uniform",
+            name="Wv",
+        )
+
+        self.Wo = self.add_weight(
+            shape=(self.d_model, self.d_model),
+            initializer="glorot_uniform",
+            name="Wo",
+        )
+
+    def call(self, x, mask=None):
+
+        batch = x.shape[0]
+        seq = x.shape[1]
+
+        Q = x @ self.Wq
+        K = x @ self.Wk
+        V = x @ self.Wv
+
+        Q = Q.reshape(
+            batch,
+            seq,
+            self.num_heads,
+            self.head_dim,
+        )
+
+        K = K.reshape(
+            batch,
+            seq,
+            self.num_heads,
+            self.head_dim,
+        )
+
+        V = V.reshape(
+            batch,
+            seq,
+            self.num_heads,
+            self.head_dim,
+        )
+
+        Q = Q.transpose(0, 2, 1, 3)
+        K = K.transpose(0, 2, 1, 3)
+        V = V.transpose(0, 2, 1, 3)
+
+        scores = Q @ K.transpose(0, 1, 3, 2)
+
+        # scores = scores / math.sqrt(self.head_dim)
+        scores = scores / Tensor(self.head_dim).sqrt()
+        # mask = np.triu(
+        #     np.ones((seq, seq)),
+        #     k=1
+        # ).astype(bool)
+
+        if mask is not None:
+            scores = scores.masked_fill(mask == 0, -1e9)
+
+        weights = scores.softmax(axis=-1)
+
+        context = weights @ V
+
+        context = context.transpose(0, 2, 1, 3)
+
+        context = context.reshape(
+            batch,
+            seq,
+            self.d_model,
+        )
+
+        output = context @ self.Wo
+
+        return output
+
+    def state_dict(self):
+        return {
+            "Wq": self.Wq.data.copy(),
+            "Wk": self.Wk.data.copy(),
+            "Wv": self.Wv.data.copy(),
+            "Wo": self.Wo.data.copy(),
+        }
+
+    def load_state_dict(self, state):
+        self.Wq.data[:] = state["Wq"]
+        self.Wk.data[:] = state["Wk"]
+        self.Wv.data[:] = state["Wv"]
+        self.Wo.data[:] = state["Wo"]

src/models/transformers/TransformerBlock.py

@@ -0,0 +1,79 @@
+from src.activations import activation_fns
+from src.models.transformers.Dropout import Dropout
+from src.models.transformers.LayerNorm import LayerNorm
+from src.models.transformers.MultiHeadAttention import MultiHeadAttention
+from src.neural.Dense import Dense
+from src.neural.Layer import Layer
+
+
+class TransformerBlock(Layer):
+
+    def __init__(self, d_model, num_heads, ff_dim):
+        super().__init__()
+        self.attn = MultiHeadAttention(d_model, num_heads)
+        self.norm1 = LayerNorm(d_model)
+
+        self.fc1 = Dense(ff_dim)
+        self.fc2 = Dense(d_model)
+
+        self.norm2 = LayerNorm(d_model)
+        self.dropout = Dropout(0.1)
+
+    def call(self, x):
+        h = self.attn(x)
+        x = self.norm1(x + self.dropout(h))
+        # x = self.norm1(x + h)
+
+        h = self.fc2(activation_fns["gelu"](self.fc1(x)))
+        x = self.norm2(x + self.dropout(h))
+        # x = self.norm2(x + h)
+
+        return x
+
+    def build(self, input_shape):
+        self.attn.build(input_shape)
+
+        self.norm1.build(input_shape)
+
+        self.fc1.build(input_shape)
+
+        ff_shape = (*input_shape[:-1], self.fc1.units)
+
+        self.fc2.build(ff_shape)
+
+        self.norm2.build(input_shape)
+
+        self.dropout.build(input_shape)
+
+        self.built = True
+
+    def parameters(self):
+        params = []
+
+        params.extend(self.attn.parameters())
+
+        params.extend(self.norm1.parameters())
+
+        params.extend(self.fc1.parameters())
+
+        params.extend(self.fc2.parameters())
+
+        params.extend(self.norm2.parameters())
+
+        return params
+
+    def state_dict(self):
+        return {
+            "attn": self.attn.state_dict(),
+            "norm1": self.norm1.state_dict(),
+            "fc1": self.fc1.state_dict(),
+            "fc2": self.fc2.state_dict(),
+            "norm2": self.norm2.state_dict(),
+        }
+
+    def load_state_dict(self, state):
+        self.attn.load_state_dict(state["attn"])
+        self.norm1.load_state_dict(state["norm1"])
+        self.fc1.load_state_dict(state["fc1"])
+        self.fc2.load_state_dict(state["fc2"])
+        self.norm2.load_state_dict(state["norm2"])

src/neural/Dense.py

@@ -0,0 +1,58 @@
+from src.neural.Layer import Layer
+from src.activations import activation_fns
+
+
+class Dense(Layer):
+    def __init__(self, units, activation=None,
+        kernel_initializer = "glorot_uniform",
+        bias_initializer = "zeros",
+        trainable = True,
+    ):
+        super().__init__()
+        self.units = units
+        self.activation = activation
+
+        self.kernel_initializer = kernel_initializer
+        self.bias_initializer = bias_initializer
+        self.trainable = trainable
+
+        self.kernel = None
+        self.bias = None
+
+    def build(self, input_shape):
+        in_features = input_shape[-1]
+
+        self.kernel = self.add_weight(
+            # row = in_features, column = units
+            shape=(in_features, self.units),
+            initializer=self.kernel_initializer,
+            trainable=self.trainable,
+            name="kernel",
+        )
+
+        self.bias = self.add_weight(
+            shape=(self.units,),
+            initializer=self.bias_initializer,
+            trainable=self.trainable,
+            name="bias",
+        )
+
+    def call(self, inputs):
+        output = inputs @ self.kernel
+        output = output + self.bias
+
+        if self.activation is not None:
+            output = activation_fns[self.activation](output)
+
+        return output
+
+    def compute_output_shape(self, input_shape):
+        return input_shape[:-1] + (self.units,)
+
+    def get_config(self):
+        return {
+            "units": self.units,
+            "activation": self.activation,
+            "kernel_initializer": self.kernel_initializer,
+            "bias_initializer": self.bias_initializer,
+        }

src/neural/LSTM.py

@@ -0,0 +1,167 @@
+from src.core.Tensor import Tensor
+from src.neural.Layer import Layer
+
+
+class LSTM(Layer):
+
+    def __init__(
+        self,
+        hidden_size,
+        return_sequences=False,
+    ):
+        super().__init__()
+
+        self.bo = None
+        self.Uo = None
+        self.Wo = None
+        self.bg = None
+        self.Ug = None
+        self.Wg = None
+        self.bf = None
+        self.Uf = None
+        self.Wf = None
+        self.bi = None
+        self.Ui = None
+        self.Wi = None
+        self.hidden_size = hidden_size
+        self.return_sequences = return_sequences
+
+    def build(self, input_shape):
+
+        input_size = input_shape[-1]
+
+        self.Wi = self.add_weight(
+            shape=(input_size, self.hidden_size),
+            initializer="glorot_uniform",
+            name="Wi",
+        )
+
+        self.Ui = self.add_weight(
+            shape=(self.hidden_size, self.hidden_size),
+            initializer="orthogonal",
+            name="Ui",
+        )
+
+        self.bi = self.add_weight(
+            shape=(self.hidden_size,),
+            initializer="zeros",
+            name="bi",
+        )
+
+        self.Wf = self.add_weight(
+            shape=(input_size, self.hidden_size),
+            initializer="glorot_uniform",
+            name="Wf",
+        )
+
+        self.Uf = self.add_weight(
+            shape=(self.hidden_size, self.hidden_size),
+            initializer="orthogonal",
+            name="Uf",
+        )
+
+        self.bf = self.add_weight(
+            shape=(self.hidden_size,),
+            initializer="zeros",
+            name="bf",
+        )
+
+        self.Wg = self.add_weight(
+            shape=(input_size, self.hidden_size),
+            initializer="glorot_uniform",
+            name="Wg",
+        )
+
+        self.Ug = self.add_weight(
+            shape=(self.hidden_size, self.hidden_size),
+            initializer="orthogonal",
+            name="Ug",
+        )
+
+        self.bg = self.add_weight(
+            shape=(self.hidden_size,),
+            initializer="zeros",
+            name="bg",
+        )
+
+        self.Wo = self.add_weight(
+            shape=(input_size, self.hidden_size),
+            initializer="glorot_uniform",
+            name="Wo",
+        )
+
+        self.Uo = self.add_weight(
+            shape=(self.hidden_size, self.hidden_size),
+            initializer="orthogonal",
+            name="Uo",
+        )
+
+        self.bo = self.add_weight(
+            shape=(self.hidden_size,),
+            initializer="zeros",
+            name="bo",
+        )
+
+    def call(self, x):
+
+        batch = x.shape[0]
+        timesteps = x.shape[1]
+
+        h = Tensor.zeros(
+            (batch, self.hidden_size),
+            requires_grad=False,
+        )
+
+        c = Tensor.zeros(
+            (batch, self.hidden_size),
+            requires_grad=False,
+        )
+
+        outputs = []
+
+        for t in range(timesteps):
+
+            xt = x[:, t, :]
+
+            i = (
+                xt @ self.Wi +
+                h @ self.Ui +
+                self.bi
+            ).sigmoid()
+
+            f = (
+                xt @ self.Wf +
+                h @ self.Uf +
+                self.bf
+            ).sigmoid()
+
+            g = (
+                xt @ self.Wg +
+                h @ self.Ug +
+                self.bg
+            ).tanh()
+
+            o = (
+                xt @ self.Wo +
+                h @ self.Uo +
+                self.bo
+            ).sigmoid()
+
+            c = f * c + i * g
+
+            h = o * c.tanh()
+
+            if self.return_sequences:
+                outputs.append(h)
+
+        if self.return_sequences:
+            return Tensor.stack(outputs, axis=1)
+
+        return h
+
+    def get_config(self):
+        return {
+            "class_name": "LSTM",
+            "hidden_size": self.hidden_size,
+            "return_sequences": self.return_sequences,
+        }

src/neural/Layer.py

@@ -0,0 +1,72 @@
+from src.neural.Parameter import Parameter
+from src.initializers import initializer_fns
+
+class Layer:
+    def __init__(self):
+        self.built = False
+
+        self.weights = []
+        self.trainable_weights = []
+        self.non_trainable_weights = []
+
+    def __call__(self, inputs):
+        if not self.built:
+            self.build(inputs.shape)
+            self.built = True
+
+        return self.call(inputs)
+
+    def build(self, input_shape):
+        raise NotImplementedError
+
+    def call(self, inputs):
+        raise NotImplementedError
+
+    def add_weight(
+        self,
+        shape,
+        initializer="glorot_uniform",
+        trainable=True,
+        name=None,
+    ):
+
+        value = initializer_fns[initializer](shape)
+
+        parameter = Parameter(
+            data=value,
+            trainable=trainable,
+            name=name,
+        )
+
+        self.weights.append(parameter)
+
+        if trainable:
+            self.trainable_weights.append(parameter)
+        else:
+            self.non_trainable_weights.append(parameter)
+
+        return parameter
+
+    def parameters(self):
+        return self.trainable_weights
+
+    def state_dict(self):
+        state = {}
+
+        for p in self.parameters():
+            state[p.name] = p.data.copy()
+
+        return state
+
+    def load_state_dict(self, state):
+
+        if not state:
+            return
+
+        for p in self.parameters():
+
+            if p.name in state:
+                p.data[:] = state[p.name]
+
+    def compute_output_shape(self, input_shape):
+        raise NotImplementedError

src/neural/Parameter.py

@@ -0,0 +1,30 @@
+import numpy as np
+
+from src.core.Tensor import Tensor
+
+
+# class Parameter:
+#     def __init__(self, data, trainable=True, name=None):
+#         self.data = np.asarray(data, dtype=np.float32)
+#         self.grad = np.zeros_like(self.data)
+#
+#         self.trainable = trainable
+#         self.name = name
+#
+#     def zero_grad(self):
+#         self.grad.fill(0)
+#
+#     @property
+#     def shape(self):
+#         return self.data.shape
+
+class Parameter(Tensor):
+
+    def __init__(self, data, trainable=True, name=None):
+        super().__init__(
+            data,
+            requires_grad=True,
+        )
+
+        self.name = name
+        self.trainable = trainable

src/neural/RNN.py

@@ -0,0 +1,83 @@
+from src.neural.Layer import Layer
+from src.activations import Tanh
+from src.core.Tensor import Tensor
+
+
+class RNN(Layer):
+
+    def __init__(
+        self,
+        hidden_size,
+        return_sequences=False,
+    ):
+        super().__init__()
+
+        self.hidden_size = hidden_size
+        self.return_sequences = return_sequences
+
+        self.Wx = None
+        self.Wh = None
+        self.bh = None
+
+    def build(self, input_shape):
+        input_size = input_shape[-1]
+
+        self.Wx = self.add_weight(
+            shape=(input_size, self.hidden_size),
+            initializer="glorot_uniform",
+            trainable=True,
+            name="Wx",
+        )
+
+        self.Wh = self.add_weight(
+            shape=(self.hidden_size, self.hidden_size),
+            initializer="orthogonal",
+            trainable=True,
+            name="Wh",
+        )
+
+        self.bh = self.add_weight(
+            shape=(self.hidden_size,),
+            initializer="zeros",
+            trainable=True,
+            name="bh",
+        )
+
+    def call(self, x):
+
+        batch_size = x.shape[0]
+        time_steps = x.shape[1]
+
+        h = Tensor.zeros(
+            (batch_size, self.hidden_size),
+            requires_grad=False,
+        )
+
+        outputs = []
+
+        for t in range(time_steps):
+
+            xt = x[:, t, :]
+
+            h = Tanh.forward(
+                xt @ self.Wx +
+                h @ self.Wh +
+                self.bh
+            )
+
+            if self.return_sequences:
+                outputs.append(h)
+
+        if self.return_sequences:
+            return Tensor.stack(outputs, axis=1)
+
+        return h
+
+    def compute_output_shape(self, input_shape):
+
+        batch = input_shape[0]
+
+        return (
+            batch,
+            self.hidden_size,
+        )

src/ops/stack.py

@@ -0,0 +1,40 @@
+import numpy as np
+
+from src.core.Tensor import Tensor
+
+
+class Stack:
+
+    @staticmethod
+    def forward(tensors, axis=0):
+
+        data = np.stack(
+            [t.data for t in tensors],
+            axis=axis,
+        )
+
+        out = Tensor(
+            data,
+            requires_grad=any(t.requires_grad for t in tensors),
+            parents=tuple(tensors),
+            op="Stack",
+        )
+
+        def _backward():
+
+            for i, tensor in enumerate(tensors):
+
+                if not tensor.requires_grad:
+                    continue
+
+                grad = np.take(
+                    out.grad,
+                    indices=i,
+                    axis=axis,
+                )
+
+                tensor.grad += grad
+
+        out._backward = _backward
+
+        return out

src/optimizers/Adagrad.py

@@ -0,0 +1,31 @@
+import numpy as np
+
+
+class Adagrad:
+
+    def __init__(self, lr=0.01, eps=1e-8):
+        self.lr = lr
+        self.eps = eps
+        self.cache = {}
+
+    def step(self, params):
+
+        for p in params:
+
+            if not p.requires_grad:
+                continue
+
+            if id(p) not in self.cache:
+                self.cache[id(p)] = np.zeros_like(p.data)
+
+            self.cache[id(p)] += p.grad ** 2
+
+            p.data -= (
+                self.lr
+                * p.grad
+                / (np.sqrt(self.cache[id(p)]) + self.eps)
+            )
+
+    def zero_grad(self, params):
+        for p in params:
+            p.zero_grad()