refactored train

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)

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)):