...
Code Block | ||||
---|---|---|---|---|
| ||||
public class SynchronousBoundedLinearRegression { private static final N_DIM = 50; private static final OutputTag<double[]> FINAL_MODEL_OUTPUT_TAG = new OutputTag<double[]>{}; public static void main(String[] args) { DataStream<double[]> initParameters = loadParameters().setParallelism(1); DataStream<Tuple2<double[], Double>> dataset = loadDataSet().setParallelism(1); int batch = 5; int epochEachBatch = 10; ResultStreams resultStreams = new BoundedIteration() .withBody(new IterationBody( @IterationInput("model") DataStream<double[]> model, @IterationInput("dataset") DataStream<Tuple2<double[], Double>> dataset ) { SingleOutputStreamOperator<double[]> parameters = model.process(new ParametersCacheFunction()); DataStream<double[]> modelUpdate = parameters.setParallelism(1) .broadcast() .connect(dataset) .coProcess(new TrainFunction()) .setParallelism(10) return new BoundedIterationDeclarationBuilder() .withFeedback("model", modelUpdate) .withOutput("final_model", parameters.getSideOut(FINAL_MODEL_OUTPUT_TAG)) .until(new TerminationCondition(null, context -> context.getRound() >= batch * epochEachBatch)) .build(); }) .build(); DataStream<double[]> finalModel = resultStreams.get("final_model"); finalModel.print(); } public static class ParametersCacheFunction extends ProcessFunction<double[], double[]> implements BoundedIterationProgressListener<double[]> { private final double[] parameters = new double[N_DIM]; public void processElement(double[] update, Context ctx, Collector<O> output) { // Suppose we have a util to add the second array to the first. ArrayUtils.addWith(parameters, update); } public void onRoundEnd(int[] round, Context context, Collector<T> collector) { collector.collect(parameters); } public void onIterationEnd(int[] round, Context context) { context.output(FINAL_MODEL_OUTPUT_TAG, parameters); } } public static class TrainFunction extends CoProcessFunction<double[], Tuple2<double[], Double>, double[]> implements BoundedIterationProgressListener<double[]> { private final List<Tuple2<double[], Double>> dataset = new ArrayList<>(); private double[] firstRoundCachedParameter; private Supplier<int[]> recordRoundQuerier; public void setCurrentRecordRoundsQuerier(Supplier<int[]> querier) { this.recordRoundQuerier = querier; } public void processElement1(double[] parameter, Context context, Collector<O> output) { int[] round = recordRoundQuerier.get(); if (round[0] == 0) { firstRoundCachedParameter = parameter; return; } calculateModelUpdate(parameter, output); } public void processElement2(Tuple2<double[], Double> trainSample, Context context, Collector<O> output) { dataset.add(trainSample) } public void onRoundEnd(int[] round, Context context, Collector<T> output) { if (round[0] == 0) { calculateModelUpdate(firstRoundCachedParameter, output); firstRoundCachedParameter = null; } } private void calculateModelUpdate(double[] parameters, Collector<O> output) { List<Tuple2<double[], Double>> samples = sample(dataset); double[] modelUpdate = new double[N_DIM]; for (Tuple2<double[], Double> record : samples) { double diff = (ArrayUtils.muladd(record.f0, parameters) - record.f1); ArrayUtils.addWith(modelUpdate, ArrayUtils.multiply(record.f0, diff)); } output.collect(modelUpdate); } } } |
If instead we want to do asynchronous training, we would need to do the following change:
- The Parameters vertex would not wait till round end to ensure received all the updates from the iteration. Instead, it would immediately output the current parameters values once it received the model update from one train subtask.
- To label the source of the update, we would like to change the input type to be Tuple2<Integer, double[]>. The Parameters vertex would output the
We omit the change to the graph
Code Block | ||
---|---|---|
| ||
public static class ParametersCacheFunction extends ProcessFunction<Tuple2<Integer, double[]>, Tuple2<Integer, double[]>>
implements BoundedIterationProgressListener<double[]> {
private final double[] parameters = new double[N_DIM];
public void processElement(Tuple2<Integer, double[]> update, Context ctx, Collector<Tuple2<Integer, double[]>> output) {
// Suppose we have a util to add the second array to the first.
ArrayUtils.addWith(parameters, update);
output.collect(new Tuple2<>(update.f0, parameters))
}
public void onIterationEnd(int[] round, Context context) {
context.output(FINAL_MODEL_OUTPUT_TAG, parameters);
}
} |
Online Training with Unbounded Iteration
Suppose now we would change the algorithm to unbounded iteration, compared to the offline, the differences is that
- The dataset is unbounded. The Train operator could not cache all the data in the first round.
- The training algorithm might change to others like FTRL. But we keep using SGD in this example since it does not affect showing the usage of the iteration.
We also start with the synchronous case. for online training, the Train vertex usually do one update after accumulating one mini-batch. This is to ensure the distribution of the samples is similar to the global statistics. In this example we omit the complex data re-sample process and just fetch the next several records as one mini-batch.
The JobGraph for online training is still shown in Figure 1, with the training dataset become unbounded. Similar to the bounded cases, for the synchronous training, the process would be expected like
- The Parameters broadcast the initialized values on received the input values.
- All the Train task read the next mini-batch of records, Calculating an update and emit to the Parameters vertex. Then it would wait till received update parameters from the Parameters Vertex before it head to process the next mini-batch.
- The Parameter vertex would wait received the updates from all the Train tasks before it broadcast the updated parameters.
The code would beFor the Parameters vertex, the
Implementation Plan
Logically all the iteration types would support both BATCH and STREAM execution mode. However, according to the algorithms' requirements, we would implement
...