org.deeplearning4j.arbiter.layers.BatchNormalizationSpace 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.layers;

import org.deeplearning4j.arbiter.optimize.api.ParameterSpace;
import org.deeplearning4j.arbiter.optimize.parameter.FixedValue;
import org.deeplearning4j.arbiter.util.CollectionUtils;
import org.deeplearning4j.nn.conf.layers.BatchNormalization;

import java.util.List;

/**
 * LayerSpace for batch normalization layers
 *
 * @author Alex Black
 */
public class BatchNormalizationSpace extends FeedForwardLayerSpace {

    protected ParameterSpace decay;
    protected ParameterSpace eps;
    protected ParameterSpace isMinibatch;
    protected ParameterSpace lockGammaBeta;
    protected ParameterSpace gamma;
    protected ParameterSpace beta;

    private BatchNormalizationSpace(Builder builder) {
        super(builder);
        this.decay = builder.decay;
        this.eps = builder.eps;
        this.isMinibatch = builder.isMinibatch;
        this.lockGammaBeta = builder.lockGammaBeta;
        this.gamma = builder.gamma;
        this.beta = builder.beta;

        this.numParameters = CollectionUtils.countUnique(collectLeaves());
    }

    @Override
    public List collectLeaves() {
        List list = super.collectLeaves();
        if (decay != null) list.addAll(decay.collectLeaves());
        if (eps != null) list.addAll(eps.collectLeaves());
        if (isMinibatch != null) list.addAll(isMinibatch.collectLeaves());
        if (lockGammaBeta != null) list.addAll(lockGammaBeta.collectLeaves());
        if (gamma != null) list.addAll(gamma.collectLeaves());
        if (beta != null) list.addAll(beta.collectLeaves());
        return list;
    }


    @Override
    public BatchNormalization getValue(double[] parameterValues) {
        BatchNormalization.Builder b = new BatchNormalization.Builder();
        setLayerOptionsBuilder(b, parameterValues);
        return b.build();
    }

    protected void setLayerOptionsBuilder(BatchNormalization.Builder builder, double[] values) {
        super.setLayerOptionsBuilder(builder, values);
        if (decay != null) builder.decay(decay.getValue(values));
        if (eps != null) builder.eps(eps.getValue(values));
        if (isMinibatch != null) builder.minibatch(isMinibatch.getValue(values));
        if (lockGammaBeta != null) builder.lockGammaBeta(lockGammaBeta.getValue(values));
        if (gamma != null) builder.gamma(gamma.getValue(values));
        if (beta != null) builder.beta(beta.getValue(values));
    }

    @Override
    public String toString() {
        return toString(", ");
    }

    @Override
    public String toString(String delim) {
        StringBuilder sb = new StringBuilder();
        sb.append("BatchNormalizationSpace(").append(super.toString(delim));
        if(decay != null) sb.append("decay: ").append(decay).append(delim);
        if(eps != null) sb.append("eps: ").append(eps).append(delim);
        if(isMinibatch != null) sb.append("isMinibatch: ").append(isMinibatch).append(delim);
        if(lockGammaBeta != null) sb.append("lockGammaBeta: ").append(lockGammaBeta).append(delim);
        if(gamma != null) sb.append("gamma: ").append(gamma).append(delim);
        if(beta != null) sb.append("beta: ").append(beta).append(delim);
        sb.append(")");
        return sb.toString();
    }

    public static class Builder extends FeedForwardLayerSpace.Builder {

        protected ParameterSpace decay;
        protected ParameterSpace eps;
        protected ParameterSpace isMinibatch;
        protected ParameterSpace lockGammaBeta;
        protected ParameterSpace gamma;
        protected ParameterSpace beta;

        public Builder minibatch(boolean minibatch) {
            return minibatch(new FixedValue<>(minibatch));
        }

        public Builder minibatch(ParameterSpace minibatch) {
            this.isMinibatch = minibatch;
            return this;
        }

        public Builder gamma(double gamma) {
            return gamma(new FixedValue<>(gamma));
        }

        public Builder gamma(ParameterSpace gamma) {
            this.gamma = gamma;
            return this;
        }

        public Builder beta(double beta) {
            return beta(new FixedValue<>(beta));
        }

        public Builder beta(ParameterSpace beta) {
            this.beta = beta;
            return this;
        }

        public Builder eps(double eps) {
            return eps(new FixedValue<>(eps));
        }

        public Builder eps(ParameterSpace eps) {
            this.eps = eps;
            return this;
        }

        public Builder decay(double decay) {
            return decay(new FixedValue(decay));
        }

        public Builder decay(ParameterSpace decay) {
            this.decay = decay;
            return this;
        }

        public Builder lockGammaBeta(boolean lockGammaBeta) {
            return lockGammaBeta(new FixedValue<>(lockGammaBeta));
        }

        public Builder lockGammaBeta(ParameterSpace lockGammaBeta) {
            this.lockGammaBeta = lockGammaBeta;
            return this;
        }

        @Override
        public BatchNormalizationSpace build() {
            return new BatchNormalizationSpace(this);
        }
    }
}