43 lines
1.3 KiB
Python
43 lines
1.3 KiB
Python
from rbm.entity import Entity
|
|
from rbm.train import cd_binary_binary, cd_binary_gaussian, cd_gaussian_binary, cd_gaussian_gaussian
|
|
from rbm.matrix import Mat, rms_error_accu, np
|
|
|
|
class LossFunction:
|
|
def __init__(self, entity_type: Entity.Type):
|
|
self.func = None
|
|
if entity_type == Entity.Type.GB_RBM:
|
|
self.func = cd_gaussian_binary
|
|
if entity_type == Entity.Type.BB_RBM:
|
|
self.func = cd_binary_binary
|
|
if entity_type == Entity.Type.GG_RBM:
|
|
self.func = cd_gaussian_gaussian
|
|
if entity_type == Entity.Type.BG_RBM:
|
|
self.func = cd_binary_gaussian
|
|
|
|
def __call__(self, entity: Entity, data: Mat):
|
|
# Contrastive divergence learning: calculate gradients
|
|
dwhv, dbv, dbh = self.func(entity, data)
|
|
|
|
return dwhv, dbv, dbh
|
|
|
|
class Optimizer:
|
|
def __init__(self, entity: Entity):
|
|
self.loss = LossFunction(entity.type)
|
|
self.entity = entity
|
|
|
|
def zero_grad(self):
|
|
self.entity.grad_zero()
|
|
|
|
def step(self, data: Mat):
|
|
params = self.entity.training_params
|
|
|
|
# Call loss function
|
|
dwhv, dbv, dbh = self.loss(self.entity, data)
|
|
|
|
# Adjust weight and biases
|
|
grad = self.entity.grad_compute(dbv, dbh, dwhv, learning_rate=params.learning_rate / data.shape[0],
|
|
momentum=params.momentum, weight_decay=params.weight_decay)
|
|
# Adjust weights
|
|
self.entity.state_adjust(grad)
|
|
|