ai.djl.training

Class LocalParameterServer

java.lang.Object

ai.djl.training.LocalParameterServer

All Implemented Interfaces:

ParameterServer, java.lang.AutoCloseable

public class LocalParameterServer
extends java.lang.Object
implements ParameterServer

LocalParameterServer is an implementation of the ParameterServer interface.

Constructor Summary

Constructors

Constructor and Description
LocalParameterServer(Optimizer optimizer) Create a new instance of LocalParameterServer for the given optimizer.

Method Summary

Modifier and Type	Method and Description
void	close()
void	init(java.lang.String parameterId, NDArray[] value) Initializes the ParameterStore for the given parameter.
void	pull(java.lang.String parameterId, NDArray[] weights, int priority) Pulls the value of a key from Parameter Server to NDArrays.
void	push(java.lang.String parameterId, NDArray[] grads, int priority) Updates values of a key in Parameter Server.

Methods inherited from class java.lang.Object

clone, equals, finalize, getClass, hashCode, notify, notifyAll, toString, wait, wait, wait

Constructor Detail

LocalParameterServer

public LocalParameterServer(Optimizer optimizer)

Create a new instance of LocalParameterServer for the given optimizer.

Parameters:

optimizer - an optimizer

Method Detail

init

public void init(java.lang.String parameterId,
                 NDArray[] value)

Initializes the ParameterStore for the given parameter.

Specified by:

init in interface ParameterServer

Parameters:

parameterId - the parameter ID

value - the values to be set for the given parameter

push

public void push(java.lang.String parameterId,
                 NDArray[] grads,
                 int priority)

Updates values of a key in Parameter Server.

Specified by:

push in interface ParameterServer

Parameters:

parameterId - the key to update

grads - the values corresponding to the key, values in array will be summed when key is updated

priority - the priority of the push operation. Higher priority push operations are likely to be executed before other push actions

pull

public void pull(java.lang.String parameterId,
                 NDArray[] weights,
                 int priority)

Pulls the value of a key from Parameter Server to NDArrays.

Specified by:

pull in interface ParameterServer

Parameters:

parameterId - the key to pull

weights - the NDArrays to store the value corresponding to the key, value will be copied to the devices of the NDArrays

priority - the priority of the push operation. Higher priority push operations are likely to be executed before other push actions

close

public void close()

Specified by:

close in interface ParameterServer

Specified by:

close in interface java.lang.AutoCloseable