- wrap numpy and cupy in matrix as np

- added type alias Mat for np.ndarray

src/rbm/cd_train.py

@@ -1,10 +1,8 @@
-import numpy as np
 from collections.abc import Callable
-
 from params import RbmParams
-from matrix import prob, sample, gaussian
+from matrix import prob, sample, gaussian, Mat, np
 
-def cd_jens(v_states: np.ndarray, params: RbmParams, v_to_ph: Callable, h_to_pv: Callable):
+def cd_jens(v_states: Mat, params: RbmParams, v_to_ph: Callable, h_to_pv: Callable):
 	v_probs = prob(v_states)
 	h_states = v_to_ph(v_states)
 	h_probs = h_states

src/rbm/entity.py

@@ -1,8 +1,7 @@
-import numpy as np
 from collections.abc import Callable
 from params import RbmParams
 from state import RbmState
-from matrix import sample, prob, rms_error_accu
+from matrix import sample, prob, rms_error_accu, Mat, np
 from status import Status
 
 class Entity:
@@ -10,7 +9,7 @@ class Entity:
 		self.state = RbmState.from_layer_params(shape)
 		self.params = params
 
-	def train(self, batch: np.ndarray, cd_func: Callable, status: Status):
+	def train(self, batch: Mat, cd_func: Callable, status: Status):
 		training_remain = batch.shape[0]
 		batch_size = min(self.params.mini_batch_size, training_remain)
 		if batch_size == 0:
@@ -65,14 +64,14 @@ class Entity:
 		status.on_change({"progress": {"value": round(progress), "unit": "%"},
 						  "err_rms_total": {"value": err_rms, "unit": ""}})
 
-	def v_to_ph(self, v: np.ndarray) -> np.ndarray:
+	def v_to_ph(self, v: Mat) -> Mat:
 		state = self.state.v_to_h(v)
 		if self.params.do_gaussian_hidden:
 			return state
 
 		return prob(state)
 
-	def h_to_pv(self, h: np.ndarray) -> np.ndarray:
+	def h_to_pv(self, h: Mat) -> Mat:
 		state = self.state.h_to_v(h)
 		if self.params.do_gaussian_visible:
 			return state
@@ -80,7 +79,7 @@ class Entity:
 		return prob(state)
 
 
-	def gibbs_v_to_h(self, v: np.ndarray) -> np.ndarray:
+	def gibbs_v_to_h(self, v: Mat) -> Mat:
 		h = self.v_to_ph(v)
 		for i in range(self.params.num_gibbs_samples-1):
 			h = self.h_to_pv(h)
@@ -88,7 +87,7 @@ class Entity:
 
 		return h
 
-	def gibbs_h_to_v(self, h: np.ndarray) -> np.ndarray:
+	def gibbs_h_to_v(self, h: Mat) -> Mat:
 		v = self.h_to_pv(h)
 		for i in range(self.params.num_gibbs_samples-1):
 			v = self.v_to_ph(v)

src/rbm/layer.py

@@ -1,9 +1,9 @@
-import numpy as np
 from params import RbmParams
 from state import RbmState
 from status import Status
 from cd_train import cd_jens
 from entity import Entity
+from matrix import Mat, np
 
 class Layer:
 	def __init__(self, name: str, shape: tuple[int, int, int, int], params: RbmParams):
@@ -43,7 +43,7 @@ def xor():
 	layer.load()
 
 	# Prepare training data
-	training_batch = np.array([[0,1,1], [0,0,0], [1,1,0], [1,0,1]], dtype=np.float64)
+	training_batch = Mat([[0,1,1], [0,0,0], [1,1,0], [1,0,1]], dtype=np.float64)
 
 	# Train layer
 	layer.entity.train(training_batch, cd_jens, Status())
@@ -52,7 +52,7 @@ def xor():
 	layer.save()
 
 	# Test with test data
-	test_batch = np.array([[0,0,0], [0,1,0], [1,0,0], [1,1,0]], dtype=np.float64)
+	test_batch = Mat([[0,0,0], [0,1,0], [1,0,0], [1,1,0]], dtype=np.float64)
 	for pattern in test_batch:
 		h = layer.entity.gibbs_v_to_h(pattern)
 		v = layer.entity.gibbs_h_to_v(h)

src/rbm/matrix.py

@@ -1,25 +1,36 @@
-import numpy as np
 import time
+from typing import TypeAlias
+
+USE_CUDA = 1
+if USE_CUDA:
+	import cupy as np
+	Mat: TypeAlias = np.array
+	def convert(src: Mat):
+		return np.asnumpy(src)
+else:
+	import numpy as np
+	Mat: TypeAlias = np.array
+	def convert(src: Mat):
+		return src
 
 np.random.seed(int(time.monotonic()))
-
-def uniform(shape: tuple, mu: float = 0.5, std: float = 1.0) -> np.ndarray:
+def uniform(shape: tuple, mu: float = 0.5, std: float = 1.0) -> Mat:
 	return std * (np.random.rand(shape[0], shape[1]) + mu - 0.5)
 
-def gaussian(shape: tuple, mu: float = 0.0, std: float = 1.0) -> np.ndarray:
+def gaussian(shape: tuple, mu: float = 0.0, std: float = 1.0) -> Mat:
 	return std * (np.random.randn(shape[0], shape[1]) + mu)
 
-def sample(src: np.ndarray) -> np.ndarray:
+def sample(src: Mat) -> Mat:
 	return (src > uniform(src.shape)).astype(float)
 
-def prob(src: np.ndarray) -> np.ndarray:
+def prob(src: Mat) -> Mat:
 	return 1.0 / (1 + np.exp(-src))
 
-def rms_error(d_err: np.ndarray):
+def rms_error(d_err: Mat):
 	d_err_squared = d_err * d_err
 	return np.sum(d_err_squared, 1) / d_err_squared[1]
 
-def rms_error_accu(d_err: np.ndarray):
+def rms_error_accu(d_err: Mat):
 	d_err_squared = d_err * d_err
 	s =  np.sum(d_err_squared) / d_err.size
 	return s

src/rbm/stack_deep.py

@@ -1,14 +1,14 @@
-import numpy as np
 from status import Status
 from cd_train import cd_jens
 from stack import Stack, StackType
+from matrix import Mat, np
 
 class StackDeep(Stack):
 	def __init__(self, name: str, work_dir: str = '.'):
 		Stack.__init__(self, StackType.Deep, name, work_dir)
 
-	def batch_from(self, batch: np.ndarray, from_layer_id: int = 0):
-		_batch = np.matrix.copy(batch)
+	def batch_from(self, batch: Mat, from_layer_id: int = 0):
+		_batch = np.copy(batch)
 		for index, layer in enumerate(self.layers):
 			if index == from_layer_id:
 				break
@@ -16,27 +16,27 @@ class StackDeep(Stack):
 
 		return _batch
 
-	def train(self, batch: np.ndarray, status=Status()):
-		_batch = np.matrix.copy(batch)
+	def train(self, batch: Mat, status=Status()):
+		_batch = np.copy(batch)
 		for index, layer in enumerate(self.layers):
 			print(f"Train layer {index} for {layer.entity.params.num_epochs} epochs")
 			_batch = self.batch_from(batch, index)
 			layer.entity.train(_batch, cd_func=cd_jens, status=status)
 
 
-	def pass_up(self, visible: np.ndarray, from_layer_id: int = 0):
-		h = np.matrix.copy(visible)
+	def pass_up(self, visible: Mat, from_layer_id: int = 0):
+		h = np.copy(visible)
 		for layer in self.layers[from_layer_id:]:
 			h = layer.entity.gibbs_v_to_h(h)
 		return h
 
-	def pass_down(self, hidden: np.ndarray, from_layer_id: int = 0):
-		v = np.matrix.copy(hidden)
+	def pass_down(self, hidden: Mat, from_layer_id: int = 0):
+		v = np.copy(hidden)
 		for layer in list(reversed(self.layers))[from_layer_id:]:
 			v = layer.entity.gibbs_h_to_v(v)
 		return v
 
-	def pass_down_up(self, visible: np.ndarray, from_layer_id: int = 0):
+	def pass_down_up(self, visible: Mat, from_layer_id: int = 0):
 		h = self.pass_up(visible, from_layer_id)
 		v = self.pass_down(h)
 		return v

src/rbm/stack_factory.py

@@ -1,13 +1,11 @@
 import json
 from collections.abc import Callable
-from stack import Stack, StackType
+from stack import StackType
 from layer import Layer
 from params import RbmParams
-
 from stack_deep import StackDeep
 from stack_rnn import StackRnn
 
-
 class StackFactory:
 	@classmethod
 	def from_dict(cls, project: dict, work_dir: str = ".", layer_constructor: Callable = None) -> StackDeep|StackRnn:

src/rbm/state.py

@@ -1,9 +1,7 @@
-import numpy as np
-from matrix import uniform
-
+from matrix import uniform, Mat, np
 
 class RbmState:
-	def __init__(self, w_hv: np.ndarray, b_v: np.ndarray, b_h: np.ndarray):
+	def __init__(self, w_hv: Mat, b_v: Mat, b_h: Mat):
 		self.num_visible, self.num_hidden = w_hv.shape
 		self.w_hv = w_hv
 		self.b_v = b_v
@@ -32,10 +30,10 @@ class RbmState:
 
 		return obj
 
-	def v_to_h(self, visible: np.ndarray) -> np.ndarray:
+	def v_to_h(self, visible: Mat) -> Mat:
 		return np.dot(visible, self.w_hv) + self.b_h
 
-	def h_to_v(self, hidden: np.ndarray) -> np.ndarray:
+	def h_to_v(self, hidden: Mat) -> Mat:
 		return np.dot(hidden, np.transpose(self.w_hv)) + self.b_v
 
 	def to_file(self, filename: str):

src/rbm/test_rbm.py

@@ -1,18 +1,18 @@
 import os.path
-import numpy as np
 import cv2 as cv
 from argparse import ArgumentParser
 from stack_factory import StackFactory
 from status import Status
 from stack_deep import StackDeep
+from matrix import Mat, np, convert
 
-def cv_show(name: str, vec: np.array, shape):
-	img = cv.Mat(np.resize(vec, shape))
+def cv_show(name: str, vec: Mat, shape):
+	img = cv.Mat(convert(np.resize(vec, shape)))
 	img_n = cv.normalize(src=img, dst=None, alpha=255, beta=0, norm_type=cv.NORM_MINMAX, dtype=cv.CV_8U)
 	cv.imshow(f"{name}", img_n)
 
 class MyStatus(Status):
-	def __init__(self, _stack: StackDeep, _batch: np.ndarray):
+	def __init__(self, _stack: StackDeep, _batch: Mat):
 		Status.__init__(self, update_interval=10)
 		self.stack = _stack
 		self.batch = _batch
@@ -41,7 +41,7 @@ class MyStatus(Status):
 
 		return do_continue
 
-def read_armadillo(filename: str) -> np.ndarray:
+def read_armadillo(filename: str) -> Mat:
 	result = None
 	with open(filename) as fp:
 		identifier = fp.readline().replace("\n", '')
@@ -57,7 +57,7 @@ def read_armadillo(filename: str) -> np.ndarray:
 			line = fp.readline().replace("\n", '').split('  ')
 			line = line[1:]
 			data = [float(s) for s in line]
-			result[row, :] = data
+			result[row, :] = Mat(data)
 
 	return result