ai.djl.training

Interface ParameterServer

All Superinterfaces:

java.lang.AutoCloseable

All Known Implementing Classes:

LocalParameterServer

public interface ParameterServer
extends java.lang.AutoCloseable

An interface for a key-value store to store parameters, and their corresponding gradients.

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.

Method Detail

init

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

Initializes the ParameterStore for the given parameter.

Parameters:

parameterId - the parameter ID

value - the values to be set for the given parameter

push

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

Updates values of a key in Parameter Server.

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

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

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

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

void close()

Specified by:

close in interface java.lang.AutoCloseable