ai.djl.training

Interface GradientCollector

All Superinterfaces:

java.lang.AutoCloseable

public interface GradientCollector
extends java.lang.AutoCloseable

An interface that provides a mechanism to collect gradients during training.

The GradientCollector should be opened with a try-with-resources. All operations performed within the try-with-resources are recorded and the variables marked. When backward function is called, gradients are collected w.r.t previously marked variables.

Method Summary

Modifier and Type	Method and Description
void	backward(NDArray target) Calculate the gradient w.r.t previously marked variable (head).
void	close()

Method Detail

backward

void backward(NDArray target)

Calculate the gradient w.r.t previously marked variable (head).

Parameters:

target - the target NDArray to calculate the gradient w.r.t head

close

void close()

Specified by:

close in interface java.lang.AutoCloseable