001/*
002 * Copyright (c) 2015-2020, Oracle and/or its affiliates. All rights reserved.
003 *
004 * Licensed under the Apache License, Version 2.0 (the "License");
005 * you may not use this file except in compliance with the License.
006 * You may obtain a copy of the License at
007 *
008 *     http://www.apache.org/licenses/LICENSE-2.0
009 *
010 * Unless required by applicable law or agreed to in writing, software
011 * distributed under the License is distributed on an "AS IS" BASIS,
012 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express implied.
013 * See the License for the specific language governing permissions and
014 * limitations under the License.
015 */
016
017package org.tribuo.classification.sgd.crf;
018
019import com.oracle.labs.mlrg.olcut.config.Config;
020import com.oracle.labs.mlrg.olcut.provenance.Provenance;
021import com.oracle.labs.mlrg.olcut.util.Pair;
022import org.tribuo.ImmutableFeatureMap;
023import org.tribuo.ImmutableOutputInfo;
024import org.tribuo.WeightedExamples;
025import org.tribuo.classification.Label;
026import org.tribuo.classification.sgd.Util;
027import org.tribuo.math.StochasticGradientOptimiser;
028import org.tribuo.math.la.SGDVector;
029import org.tribuo.math.la.Tensor;
030import org.tribuo.provenance.ModelProvenance;
031import org.tribuo.provenance.TrainerProvenance;
032import org.tribuo.provenance.impl.TrainerProvenanceImpl;
033import org.tribuo.sequence.SequenceDataset;
034import org.tribuo.sequence.SequenceExample;
035import org.tribuo.sequence.SequenceTrainer;
036
037import java.time.OffsetDateTime;
038import java.util.Map;
039import java.util.SplittableRandom;
040import java.util.logging.Logger;
041
042/**
043 * A trainer for CRFs using SGD. Modelled after FACTORIE's trainer for CRFs.
044 * <p>
045 * See:
046 * <pre>
047 * Lafferty J, McCallum A, Pereira FC.
048 * "Conditional Random Fields: Probabilistic Models for Segmenting and Labeling Sequence Data"
049 * Proceedings of the 18th International Conference on Machine Learning 2001 (ICML 2001).
050 * </pre>
051 */
052public class CRFTrainer implements SequenceTrainer<Label>, WeightedExamples {
053    private static final Logger logger = Logger.getLogger(CRFTrainer.class.getName());
054
055    @Config(mandatory = true,description="The gradient optimiser to use.")
056    private StochasticGradientOptimiser optimiser;
057
058    @Config(description="The number of gradient descent epochs.")
059    private int epochs = 5;
060
061    @Config(description="Log values after this many updates.")
062    private int loggingInterval = -1;
063
064    @Config(description="Minibatch size in SGD.")
065    private int minibatchSize = 1;
066
067    @Config(mandatory = true,description="Seed for the RNG used to shuffle elements.")
068    private long seed;
069
070    @Config(description="Shuffle the data before each epoch. Only turn off for debugging.")
071    private boolean shuffle = true;
072
073    private SplittableRandom rng;
074
075    private int trainInvocationCounter;
076
077    /**
078     * Creates a CRFTrainer which uses SGD to learn the parameters.
079     * @param optimiser The gradient optimiser to use.
080     * @param epochs The number of SGD epochs (complete passes through the training data).
081     * @param loggingInterval Log the loss after this many iterations. If -1 don't log anything.
082     * @param minibatchSize The size of the minibatches used to aggregate gradients.
083     * @param seed A seed for the random number generator, used to shuffle the examples before each epoch.
084     */
085    public CRFTrainer(StochasticGradientOptimiser optimiser, int epochs, int loggingInterval, int minibatchSize, long seed) {
086        this.optimiser = optimiser;
087        this.epochs = epochs;
088        this.loggingInterval = loggingInterval;
089        this.minibatchSize = minibatchSize;
090        this.seed = seed;
091        postConfig();
092    }
093
094    /**
095     * Sets the minibatch size to 1.
096     * @param optimiser The gradient optimiser to use.
097     * @param epochs The number of SGD epochs (complete passes through the training data).
098     * @param loggingInterval Log the loss after this many iterations. If -1 don't log anything.
099     * @param seed A seed for the random number generator, used to shuffle the examples before each epoch.
100     */
101    public CRFTrainer(StochasticGradientOptimiser optimiser, int epochs, int loggingInterval, long seed) {
102        this(optimiser,epochs,loggingInterval,1,seed);
103    }
104
105    /**
106     * Sets the minibatch size to 1 and the logging interval to 100.
107     * @param optimiser The gradient optimiser to use.
108     * @param epochs The number of SGD epochs (complete passes through the training data).
109     * @param seed A seed for the random number generator, used to shuffle the examples before each epoch.
110     */
111    public CRFTrainer(StochasticGradientOptimiser optimiser, int epochs, long seed) {
112        this(optimiser,epochs,100,1,seed);
113    }
114
115    /**
116     * For olcut.
117     */
118    private CRFTrainer() { }
119
120    /**
121     * Used by the OLCUT configuration system, and should not be called by external code.
122     */
123    @Override
124    public synchronized void postConfig() {
125        this.rng = new SplittableRandom(seed);
126    }
127
128    /**
129     * Turn on or off shuffling of examples.
130     * <p>
131     * This isn't exposed in the constructor as it defaults to on.
132     * This method should be used for debugging.
133     * @param shuffle If true shuffle the examples, if false leave them in their current order.
134     */
135    public void setShuffle(boolean shuffle) {
136        this.shuffle = shuffle;
137    }
138
139    @Override
140    public CRFModel train(SequenceDataset<Label> sequenceExamples, Map<String, Provenance> runProvenance) {
141        if (sequenceExamples.getOutputInfo().getUnknownCount() > 0) {
142            throw new IllegalArgumentException("The supplied Dataset contained unknown Outputs, and this Trainer is supervised.");
143        }
144        // Creates a new RNG, adds one to the invocation count, generates a local optimiser.
145        SplittableRandom localRNG;
146        TrainerProvenance trainerProvenance;
147        StochasticGradientOptimiser localOptimiser;
148        synchronized(this) {
149            localRNG = rng.split();
150            localOptimiser = optimiser.copy();
151            trainerProvenance = getProvenance();
152            trainInvocationCounter++;
153        }
154        ImmutableOutputInfo<Label> labelIDMap = sequenceExamples.getOutputIDInfo();
155        ImmutableFeatureMap featureIDMap = sequenceExamples.getFeatureIDMap();
156        SGDVector[][] sgdFeatures = new SGDVector[sequenceExamples.size()][];
157        int[][] sgdLabels = new int[sequenceExamples.size()][];
158        double[] weights = new double[sequenceExamples.size()];
159        int n = 0;
160        for (SequenceExample<Label> example : sequenceExamples) {
161            weights[n] = example.getWeight();
162            Pair<int[],SGDVector[]> pair = CRFModel.convertToVector(example,featureIDMap,labelIDMap);
163            sgdFeatures[n] = pair.getB();
164            sgdLabels[n] = pair.getA();
165            n++;
166        }
167        logger.info(String.format("Training SGD CRF with %d examples", n));
168
169        CRFParameters crfParameters = new CRFParameters(featureIDMap.size(),labelIDMap.size());
170
171        localOptimiser.initialise(crfParameters);
172        double loss = 0.0;
173        int iteration = 0;
174
175        for (int i = 0; i < epochs; i++) {
176            if (shuffle) {
177                Util.shuffleInPlace(sgdFeatures, sgdLabels, weights, localRNG);
178            }
179            if (minibatchSize == 1) {
180                /*
181                 * Special case a minibatch of size 1. Directly updates the parameters after each
182                 * example rather than aggregating.
183                 */
184                for (int j = 0; j < sgdFeatures.length; j++) {
185                    Pair<Double,Tensor[]> output = crfParameters.valueAndGradient(sgdFeatures[j],sgdLabels[j]);
186                    loss += output.getA()*weights[j];
187
188                    //Update the gradient with the current learning rates
189                    Tensor[] updates = localOptimiser.step(output.getB(),weights[j]);
190
191                    //Apply the update to the current parameters.
192                    crfParameters.update(updates);
193
194                    iteration++;
195                    if ((iteration % loggingInterval == 0) && (loggingInterval != -1)) {
196                        logger.info("At iteration " + iteration + ", average loss = " + loss/loggingInterval);
197                        loss = 0.0;
198                    }
199                }
200            } else {
201                Tensor[][] gradients = new Tensor[minibatchSize][];
202                for (int j = 0; j < sgdFeatures.length; j += minibatchSize) {
203                    double tempWeight = 0.0;
204                    int curSize = 0;
205                    //Aggregate the gradient updates for each example in the minibatch
206                    for (int k = j; k < j+minibatchSize && k < sgdFeatures.length; k++) {
207                        Pair<Double,Tensor[]> output = crfParameters.valueAndGradient(sgdFeatures[j],sgdLabels[j]);
208                        loss += output.getA()*weights[k];
209                        tempWeight += weights[k];
210
211                        gradients[k-j] = output.getB();
212                        curSize++;
213                    }
214                    //Merge the values into a single gradient update
215                    Tensor[] updates = crfParameters.merge(gradients,curSize);
216                    for (Tensor update : updates) {
217                        update.scaleInPlace(minibatchSize);
218                    }
219                    tempWeight /= minibatchSize;
220                    //Update the gradient with the current learning rates
221                    updates = localOptimiser.step(updates,tempWeight);
222                    //Apply the gradient.
223                    crfParameters.update(updates);
224
225                    iteration++;
226                    if ((loggingInterval != -1) && (iteration % loggingInterval == 0)) {
227                        logger.info("At iteration " + iteration + ", average loss = " + loss/loggingInterval);
228                        loss = 0.0;
229                    }
230                }
231            }
232        }
233        localOptimiser.finalise();
234        //public CRFModel(String name, String description, ImmutableInfoMap featureIDMap, ImmutableInfoMap outputIDInfo, CRFParameters parameters) {
235        ModelProvenance provenance = new ModelProvenance(CRFModel.class.getName(),OffsetDateTime.now(),sequenceExamples.getProvenance(),trainerProvenance,runProvenance);
236        CRFModel model = new CRFModel("crf-sgd-model",provenance,featureIDMap,labelIDMap,crfParameters);
237        localOptimiser.reset();
238        return model;
239    }
240
241    @Override
242    public int getInvocationCount() {
243        return trainInvocationCounter;
244    }
245
246    @Override
247    public String toString() {
248        return "CRFTrainer(optimiser="+optimiser.toString()+",epochs="+epochs+",minibatchSize="+minibatchSize+",seed="+seed+")";
249    }
250
251    @Override
252    public TrainerProvenance getProvenance() {
253        return new TrainerProvenanceImpl(this);
254    }
255}