Status
Current state: <TODO> "Under Discussion"
Discussion thread: <TODO>
...
Please keep the discussion on the mailing list rather than commenting on the wiki (wiki discussions get unwieldy fast).
Table of Contents |
---|
[This FLIP proposal is a joint work between Yun Gao, Dong Lin and Zhipeng Zhang]
Motivation
Iteration is a basic building block for a ML library. It is required for training ML models for both offline and online Training. In general, two types of iterations is required:
...
Code Block | ||||
---|---|---|---|---|
| ||||
public class SynchronousBoundedLinearRegression { private static final int N_DIM = 50; private static final int N_EPOCH = 5; private static final int N_BATCH_PER_EPOCH = 10; 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); DataStreamList resultStreams = IterationUtils.iterateBoundedStreamsUntilTermination(DataStreamList.of(initParameters), DataStreamList.of(dataset), (variableStreams, dataStreams) -> { DataStream<double>DataStream<double[]> parameterUpdates = variableStreams.get(0); DataStream<Tuple2<double[], Double>> dataset = dataStreams.get(0); SingleOutputStreamOperator<double[]> parameters = parameterUpdates.process(new ParametersCacheFunction()); DataStream<double[]> modelUpdate = parameters.setParallelism(1) .broadcast() .connect(dataset) .coProcess(new TrainFunction()) .setParallelism(10) return new IterationBodyResult(DataStreamList.of(modelUpdate), DataStreamList.of(parameters.getSideOut(FINAL_MODEL_OUTPUT_TAG))); }); DataStream<double[]> finalModel = resultStreams.get("final_model"); finalModel.print(); } public static class ParametersCacheFunction extends ProcessFunction<double[], double[]> implements IterationListener<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); } void onEpochWatermarkIncremented(int epochWatermark, Context context, Collector<T> collector) { if (epochWatermark < N_EPOCH * N_BATCH_PER_EPOCH) { 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 IterationListener<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) } void onEpochWatermarkIncremented(int epochWatermark, Context context, Collector<T> collector) { if (epochWatermark == 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); } } } |
Another method to implement the synchronous train on the bounded dataset is to use iterateAndReplayBoundedStreamsUntilTermination.
If instead we want to do asynchronous training, we would need to do the following change:
...