refactored

src/rbm/train.py

@@ -3,7 +3,35 @@ from .matrix import gaussian, sample, prob, rms_error_accu, Mat, np
 from .entity import Entity
 from .status import Status
 
-def cd_jens(entity: Entity, v_states: Mat):
+class TrainingParams:
+	def __init__(self):
+		# Training parameters
+		self.learning_rate = 0.1
+		self.momentum = 0.5
+		self.weight_decay = 0
+		self.num_epochs = 1000
+		self.mini_batch_size = 0
+		self.do_rao_blackwell = False
+		self.do_gibbs_sample_visible = False
+		self.do_gibbs_sample_hidden = False
+		self.do_batch_sample = False
+
+	@classmethod
+	def from_dict(cls, params: dict):
+		obj = TrainingParams()
+		obj.learning_rate = params["learningRate"]
+		obj.momentum = params["momentum"]
+		obj.weight_decay = params["weightDecay"]
+		obj.num_epochs = params["numEpochs"]
+		obj.mini_batch_size = params["miniBatchSize"]
+		obj.do_rao_blackwell = params["doRaoBlackwell"]
+		obj.do_gibbs_sample_visible = params["gibbsDoSampleVisible"]
+		obj.do_gibbs_sample_hidden = params["gibbsDoSampleHidden"]
+		obj.do_batch_sample = params["doSampleBatch"]
+
+		return obj
+
+def cd_jens(entity: Entity, v_states: Mat, params: TrainingParams):
 	v_probs = prob(v_states)
 	h_states = entity.v_to_ph(v_states)
 	h_probs = h_states
@@ -12,7 +40,7 @@ def cd_jens(entity: Entity, v_states: Mat):
 		h_states += gaussian(h_states.shape)
 	else:
 		h_probs = entity.v_to_ph(v_states)
-		if entity.params.do_rao_blackwell:
+		if params.do_rao_blackwell:
 			h_states = h_probs
 		else:
 			h_states = sample(h_probs)
@@ -24,13 +52,13 @@ def cd_jens(entity: Entity, v_states: Mat):
 
 	# Gibbs	sampling with training params
 	for i in range(entity.params.num_gibbs_samples):
-		if entity.params.do_gibbs_sample_hidden:
+		if params.do_gibbs_sample_hidden:
 			v_probs = entity.h_to_pv(sample(h_probs))
 		else:
 			v_probs = entity.h_to_pv(h_probs)
 
 		# Create hidden representation given v
-		if entity.params.do_gibbs_sample_visible:
+		if params.do_gibbs_sample_visible:
 			h_probs = entity.v_to_ph(sample(v_probs))
 		else:
 			h_probs = entity.v_to_ph(v_probs)
@@ -42,14 +70,14 @@ def cd_jens(entity: Entity, v_states: Mat):
 
 	return dw, dbv, dbh
 
-def train(entity: Entity, batch: Mat, status: Status, cd_func: Callable = cd_jens):
+def train(entity: Entity, batch: Mat, params: TrainingParams, status: Status, cd_func: Callable = cd_jens):
 	training_remain = batch.shape[0]
-	batch_size = min(entity.params.mini_batch_size, training_remain)
+	batch_size = min(params.mini_batch_size, training_remain)
 	if batch_size == 0:
 		batch_size = training_remain
 
 	status.on_change()
-	d_progress = 100.0 / (training_remain*entity.params.num_epochs)
+	d_progress = 100.0 / (training_remain*params.num_epochs)
 	progress = 0
 	batch_row_index = 0
 	keep_running = True
@@ -64,18 +92,18 @@ def train(entity: Entity, batch: Mat, status: Status, cd_func: Callable = cd_jen
 		inc_whv = np.zeros(entity.state.w_hv.shape)
 
 		v_states = mini_batch
-		if entity.params.do_batch_sample:
+		if params.do_batch_sample:
 			v_states = sample(mini_batch)
 
-		for epochs in range(entity.params.num_epochs):
+		for epochs in range(params.num_epochs):
 			# Contrastive divergence learning: calculate gradients
-			dwhv, dbv, dbh = cd_func(entity, v_states)
+			dwhv, dbv, dbh = cd_func(entity, v_states, params)
 
 			# Adjust weight and biases
-			kl = entity.params.learning_rate/batch_size
-			inc_bv = entity.params.momentum*inc_bv + kl*dbv
-			inc_bh = entity.params.momentum*inc_bh + kl*dbh
-			inc_whv = entity.params.momentum*inc_whv + kl*dwhv - entity.params.weight_decay*entity.state.w_hv
+			kl = params.learning_rate/batch_size
+			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