lib/src/main/java/actor/ParameterServerShard.java
<reponame>KaggleBusters/Store-Sales---Time-Series-Forecasting<filename>lib/src/main/java/actor/ParameterServerShard.java<gh_stars>0package actor;

import akka.actor.AbstractActor;
import utility.NNOperationTypes;

import java.io.Serializable;

import org.la4j.Matrix;
import org.la4j.matrix.dense.Basic2DMatrix;

public class ParameterServerShard extends AbstractActor implements Serializable { 
	// Receives learningRate (for weights update) and initial random weights.

	private int ps_id;
	private double learningRate;
	private Matrix weights;
	
	public ParameterServerShard(int ps_id, double learningRate, Matrix weights) {
		System.out.println("########!!!!!!" + weights);
		this.ps_id = ps_id;
		this.learningRate = learningRate;
		this.weights = weights;
	}

	@Override
	public Receive createReceive() {
		System.out.println("ParameterServerShard actor received message");
		return receiveBuilder()
				.match(NNOperationTypes.ParameterRequest.class, this::getLatestParameters)
				.match(NNOperationTypes.Gradient.class, this::updateWeights)
				.match(String.class, this::setWeights)
				.build();
	}
	
	// Just for testing framework correctness
	public void setWeights(String s) {
		this.weights = Matrix.fromCSV(s);
		System.out.println("@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@ " + weights);
		sender().tell("success", self());
	}

	public void getLatestParameters(NNOperationTypes.ParameterRequest paramReq) {
//		System.out.println("Lastest weights are: " + this.weights);
		sender().tell(this.weights.toCSV(), getSelf());
	}
	
	public void updateWeights(NNOperationTypes.Gradient g) {
		Basic2DMatrix grad = Basic2DMatrix.fromCSV(g.getGradient());
		weights = regularize(weights.add(grad.multiply(learningRate)));
	}

	public Matrix regularize(Matrix weights) {
		return weights.divide(weights.manhattanNorm());
	}
}