org.deeplearning4j.arbiter.conf.updater.NesterovsSpace Maven / Gradle / Ivy

Go to download
Show more of this group Show more artifacts with this name
Show all versions of arbiter-deeplearning4j Show documentation
There is a newer version: 1.0.0-beta7
Show newest version
package org.deeplearning4j.arbiter.conf.updater;

import lombok.Data;
import lombok.EqualsAndHashCode;
import org.deeplearning4j.arbiter.optimize.api.ParameterSpace;
import org.deeplearning4j.arbiter.optimize.parameter.FixedValue;
import org.nd4j.linalg.learning.config.IUpdater;
import org.nd4j.linalg.learning.config.Nesterovs;
import org.nd4j.linalg.schedule.ISchedule;
import org.nd4j.shade.jackson.annotation.JsonProperty;

@Data
@EqualsAndHashCode(callSuper = false)
public class NesterovsSpace extends BaseUpdaterSpace {

    protected ParameterSpace learningRate;
    protected ParameterSpace learningRateSchedule;
    protected ParameterSpace momentum;
    protected ParameterSpace momentumSchedule;

    public NesterovsSpace(ParameterSpace learningRate) {
        this(learningRate, null);
    }

    public NesterovsSpace(ParameterSpace learningRate, ParameterSpace momentum) {
        this(learningRate, null, momentum, null);
    }

    public NesterovsSpace(@JsonProperty("learningRate") ParameterSpace learningRate,
                          @JsonProperty("learningRateSchedule") ParameterSpace learningRateSchedule,
                          @JsonProperty("momentum") ParameterSpace momentum,
                          @JsonProperty("learningRateSchedule") ParameterSpace momentumSchedule){
        this.learningRate = learningRate;
        this.learningRateSchedule = learningRateSchedule;
        this.momentum = momentum;
        this.momentumSchedule = momentumSchedule;
    }

    public static NesterovsSpace withLR(ParameterSpace lr){
        return new NesterovsSpace(lr, null, null, null);
    }

    public static NesterovsSpace withLR(ParameterSpace lr, double momentum){
        return new NesterovsSpace(lr, null, new FixedValue<>(momentum), null);
    }

    public static NesterovsSpace withLR(ParameterSpace lr, ParameterSpace momentum){
        return new NesterovsSpace(lr, null, momentum, null);
    }

    public static NesterovsSpace withLRSchedule(ParameterSpace lrSchedule){
        return new NesterovsSpace(null, lrSchedule, null, null);
    }

    public static NesterovsSpace withLRSchedule(ParameterSpace lrSchedule, double momentum){
        return new NesterovsSpace(null, lrSchedule, new FixedValue<>(momentum), null);
    }

    public static NesterovsSpace withLRSchedule(ParameterSpace lrSchedule, ParameterSpace momentum){
        return new NesterovsSpace(null, lrSchedule, momentum, null);
    }


    @Override
    public IUpdater getValue(double[] parameterValues) {
        double lr = learningRate == null ? Nesterovs.DEFAULT_NESTEROV_LEARNING_RATE : learningRate.getValue(parameterValues);
        ISchedule lrS = learningRateSchedule == null ? null : learningRateSchedule.getValue(parameterValues);
        double m = momentum == null ? Nesterovs.DEFAULT_NESTEROV_MOMENTUM : momentum.getValue(parameterValues);
        ISchedule mS = momentumSchedule == null ? null : momentumSchedule.getValue(parameterValues);
        if(lrS == null){
            if(momentumSchedule == null){
                return new Nesterovs(lr, m);
            } else {
                return new Nesterovs(lr, mS);
            }
        } else {
            if(momentumSchedule == null){
                return new Nesterovs(lrS, m);
            } else {
                return new Nesterovs(lrS, mS);
            }
        }
    }
}