153 lines
4.4 KiB
Python
153 lines
4.4 KiB
Python
from .state import RbmState
|
|
from .matrix import prob, Mat, np
|
|
from enum import Enum
|
|
|
|
class TrainingParams:
|
|
def __init__(self,
|
|
learning_rate: float = 0.1,
|
|
momentum: float = 0.5,
|
|
weight_decay: float = 0.0,
|
|
l2_lambda: float = 0.0,
|
|
num_epochs: int = 1000,
|
|
num_gibbs_samples: int = 1,
|
|
mini_batch_size: int = 0,
|
|
do_rao_blackwell: bool = False,
|
|
do_gibbs_sample_visible: bool = False,
|
|
do_gibbs_sample_hidden: bool = False,
|
|
do_batch_sample: bool = False
|
|
):
|
|
# Training parameters
|
|
self.learning_rate = learning_rate
|
|
self.momentum = momentum
|
|
self.weight_decay = weight_decay
|
|
self.l2_lambda = l2_lambda
|
|
self.num_epochs = num_epochs
|
|
self.num_gibbs_samples = num_gibbs_samples
|
|
self.mini_batch_size = mini_batch_size
|
|
self.do_rao_blackwell = do_rao_blackwell
|
|
self.do_gibbs_sample_visible = do_gibbs_sample_visible
|
|
self.do_gibbs_sample_hidden = do_gibbs_sample_hidden
|
|
self.do_batch_sample = do_batch_sample
|
|
|
|
@classmethod
|
|
def from_dict(cls, params: dict):
|
|
obj = TrainingParams()
|
|
for key, value in params.items():
|
|
setattr(obj, key, value)
|
|
|
|
return obj
|
|
|
|
class EntityParams:
|
|
def __init__(self, do_gaussian_visible: bool = False, do_gaussian_hidden: bool = False, num_gibbs_samples: int = 1):
|
|
# Entity parameters
|
|
self.do_gaussian_visible = do_gaussian_visible
|
|
self.do_gaussian_hidden = do_gaussian_hidden
|
|
self.num_gibbs_samples = num_gibbs_samples
|
|
|
|
@classmethod
|
|
def from_dict(cls, params: dict):
|
|
obj = EntityParams()
|
|
if "doGaussianVisible" in params:
|
|
obj.do_gaussian_visible = params["doGaussianVisible"]
|
|
if "doGaussianHidden" in params:
|
|
obj.do_gaussian_hidden = params["doGaussianHidden"]
|
|
if "num_gibbs_samples" in params:
|
|
obj.num_gibbs_samples = params["num_gibbs_samples"]
|
|
|
|
return obj
|
|
|
|
class Entity:
|
|
class Type(Enum):
|
|
BB_RBM = "BB-RBM"
|
|
BG_RBM = "BG-RBM"
|
|
GB_RBM = "GB-RBM"
|
|
GG_RBM = "GG-RBM"
|
|
|
|
def __init__(self, shape: tuple[int, int], params: EntityParams, training_params: TrainingParams|None = None, enable_training: bool = True):
|
|
self.shape = shape
|
|
self.params = params
|
|
self.training_params = training_params
|
|
self.enable_training = enable_training
|
|
self.state = RbmState.from_layer_params(shape)
|
|
self.grad = RbmState.from_layer_params(shape)
|
|
self.name = f"Entity-{shape[0]}x{shape[1]}"
|
|
self.type = None
|
|
if params.do_gaussian_visible:
|
|
if params.do_gaussian_hidden:
|
|
self.type = Entity.Type.GG_RBM
|
|
else:
|
|
self.type = Entity.Type.GB_RBM
|
|
else:
|
|
if params.do_gaussian_hidden:
|
|
self.type = Entity.Type.BG_RBM
|
|
else:
|
|
self.type = Entity.Type.BB_RBM
|
|
|
|
|
|
def __call__(self, x: Mat):
|
|
return self.forward(x)
|
|
|
|
def grad_zero(self):
|
|
self.grad = RbmState.from_layer_params(self.shape)
|
|
|
|
def grad_compute(self, d_bv: Mat, d_bh: Mat, d_whv: Mat, learning_rate: float, momentum: float, weight_decay: float, l2_lambda: float):
|
|
# Compute gradient
|
|
self.grad.b_v = momentum * self.grad.b_v + learning_rate * d_bv
|
|
self.grad.b_h = momentum * self.grad.b_h + learning_rate * d_bh
|
|
# compute L2 term and penalize cost function (d_whv)
|
|
l2_norm = 0.5*np.sum(np.square(self.state.w_hv))
|
|
l2_term = l2_lambda*l2_norm*self.state.w_hv
|
|
self.grad.w_hv = momentum * self.grad.w_hv + learning_rate * (d_whv-l2_term) - learning_rate * weight_decay * self.state.w_hv
|
|
|
|
return self.grad
|
|
|
|
def state_adjust(self, grad: RbmState):
|
|
# Adjust state
|
|
self.state.b_v += grad.b_v
|
|
self.state.b_h += grad.b_h
|
|
self.state.w_hv += grad.w_hv
|
|
|
|
def forward(self, v: Mat, num_gibbs: int = 0) -> Mat:
|
|
num_gibbs = self.params.num_gibbs_samples if num_gibbs == 0 else num_gibbs
|
|
h = self._v_to_ph(v)
|
|
for i in range(num_gibbs-1):
|
|
h = self._h_to_pv(h)
|
|
h = self._v_to_ph(h)
|
|
|
|
return h
|
|
|
|
def reconstruct(self, h: Mat, num_gibbs: int = 0) -> Mat:
|
|
num_gibbs = self.params.num_gibbs_samples if num_gibbs == 0 else num_gibbs
|
|
v = self._h_to_pv(h)
|
|
for i in range(num_gibbs-1):
|
|
v = self._v_to_ph(v)
|
|
v = self._h_to_pv(v)
|
|
|
|
return v
|
|
|
|
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: Mat) -> Mat:
|
|
state = self.state.h_to_v(h)
|
|
if self.params.do_gaussian_visible:
|
|
return state
|
|
|
|
return prob(state)
|
|
|
|
def h_given_v(self, v: Mat) -> Mat:
|
|
state = self.state.v_to_h(v)
|
|
return state
|
|
|
|
def v_given_h(self, h: Mat) -> Mat:
|
|
state = self.state.h_to_v(h)
|
|
return state
|
|
|
|
if __name__ == "__main__":
|
|
print("Test: [passed]")
|