diff --git a/src/rbm/entity.py b/src/rbm/entity.py index 4ed8b8e..78da2a6 100644 --- a/src/rbm/entity.py +++ b/src/rbm/entity.py @@ -22,6 +22,21 @@ class Entity: self.shape = shape self.params = params self.state = RbmState.from_layer_params(shape) + self.delta_state = RbmState.from_layer_params(shape) + + def prepare(self): + self.delta_state = RbmState.from_layer_params(self.shape) + + def adjust(self, d_bv: Mat, d_bh: Mat, d_whv: Mat, learning_rate: float, momentum: float, weight_decay: float): + # Create delta + self.delta_state.b_v = (momentum * self.delta_state.b_v + learning_rate*d_bv) + self.delta_state.b_h = (momentum * self.delta_state.b_h + learning_rate*d_bh) + self.delta_state.w_hv = (momentum * self.delta_state.w_hv + learning_rate*d_whv - weight_decay*self.state.w_hv) + + # Create Adjust + self.state.b_v += self.delta_state.b_v + self.state.b_h += self.delta_state.b_h + self.state.w_hv += self.delta_state.w_hv def forward(self, v: Mat, num_gibbs: int = 1) -> Mat: h = self._v_to_ph(v) diff --git a/src/rbm/train.py b/src/rbm/train.py index 671cd4c..83a8a90 100644 --- a/src/rbm/train.py +++ b/src/rbm/train.py @@ -94,23 +94,13 @@ def train(entity: Entity, batch: Mat, params: TrainingParams, status: Status, cd if not keep_running: break - inc_bv = np.zeros(entity.state.b_v.shape) - inc_bh = np.zeros(entity.state.b_h.shape) - inc_whv = np.zeros(entity.state.w_hv.shape) - + entity.prepare() for epochs in range(params.num_epochs): # Contrastive divergence learning: calculate gradients dwhv, dbv, dbh, _ = cd_func(entity, mini_batch, params) # Adjust weight and biases - kl = params.learning_rate/batch.shape[0] - inc_bv = params.momentum*inc_bv + kl*dbv - inc_bh = params.momentum*inc_bh + kl*dbh - inc_whv = params.momentum*inc_whv + kl*dwhv - params.weight_decay*entity.state.w_hv - - entity.state.b_v += inc_bv - entity.state.b_h += inc_bh - entity.state.w_hv += inc_whv + entity.adjust(dbv, dbh, dwhv, learning_rate=params.learning_rate/batch.shape[0], momentum=params.momentum, weight_decay=params.weight_decay) # check if status update is needed if status.want_report(round(progress)):