ai.djl.mxnet.engine

Class MxParameterServer

java.lang.Object

ai.djl.mxnet.jna.NativeResource

ai.djl.mxnet.engine.MxParameterServer

All Implemented Interfaces:

ai.djl.training.ParameterServer, java.lang.AutoCloseable

public class MxParameterServer
extends NativeResource
implements ai.djl.training.ParameterServer

MxParameterServer is the MXNet implementation of ParameterServer.

Field Summary

Fields inherited from class ai.djl.mxnet.jna.NativeResource

handle

Constructor Summary

Constructors

Constructor and Description
MxParameterServer(ai.djl.training.optimizer.Optimizer optimizer) Constructs a new MxParameterServer.

Method Summary

Modifier and Type	Method and Description
void	close()
void	init(java.lang.String parameterId, ai.djl.ndarray.NDArray[] values)
void	pull(java.lang.String parameterId, ai.djl.ndarray.NDArray[] weights, int priority)
void	push(java.lang.String parameterId, ai.djl.ndarray.NDArray[] grads, int priority)

Methods inherited from class ai.djl.mxnet.jna.NativeResource

finalize, getHandle, getUid, isReleased

Methods inherited from class java.lang.Object

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

Constructor Detail

MxParameterServer

public MxParameterServer(ai.djl.training.optimizer.Optimizer optimizer)

Constructs a new MxParameterServer.

Parameters:

optimizer - the optimizer to use for the parameter server updates

Method Detail

init

public void init(java.lang.String parameterId,
                 ai.djl.ndarray.NDArray[] values)

Specified by:

init in interface ai.djl.training.ParameterServer

push

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

Specified by:

push in interface ai.djl.training.ParameterServer

pull

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

Specified by:

pull in interface ai.djl.training.ParameterServer

close

public void close()

Specified by:

close in interface ai.djl.training.ParameterServer

Specified by:

close in interface java.lang.AutoCloseable

Overrides:

close in class NativeResource