THIS IS A TEST INSTANCE. ALL YOUR CHANGES WILL BE LOST!!!!
...
The job graph of the algorithm could be shown in the Figure 3: in each round, we use the latest parameters to calculate the update to the parameters: ΔA ΔA = ∑(Y - XA)X. To achieve this, the Parameters vertex would broadcast the latest parameters to the Train vertex. Each subtask of the Train vertex holds a part of dataset. Follow the sprite of SGD, it would sample a small batch of training records, and calculate the update with the above equation. Then the Train vertex emit ΔA the Reducer would merge ΔA from all the train tasks and emit the reduced value to the Parameters node to update the parameters. The Reducer is required since we required the feedback streams have the same parallelism with the corresponding initialized streams.
| draw.io Diagram | ||||||||||||||||||||
|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|
|
...
| 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),
ReplayableDataStreamList.of(noReplay(dataset)),
IterationConfig.newBuilder().setOperatorRoundMode(ALL_ROUND).build();
(variableStreams, dataStreams) -> {
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)
.process(new ReduceFunction())
.setParallelism(1)
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);
}
}
} |
If instead we want to do asynchronous training, we would need to do the following change:
public static class ReduceFunction {
private double[] mergedValue = ArrayUtils.newArray(N_DIM);
public void processElement(double[] parameter, Context context, Collector<O> output) {
mergedValue = ArrayUtils.add(mergedValue, parameter);
}
void onEpochWatermarkIncremented(int epochWatermark, Context context, Collector<T> collector) {
collector.collect(mergedValue);
mergedValue = ArrayUtils.newArray(N_DIM);
}
}
} |
Here we use a specialized reduce function that
If instead we want to do asynchronous training, we would need to do the following change:
- The Parameters vertex would not 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 would only output the new parameters values to the Train task that send the update.
- The Reducer would not merge the received updates any more, instead, it would directly redirect the received updates.
We omit the change to the graph building code since the change is trivial (change the output type and the partitioner to be customized one). The change to the Parameters vertex is the follows:
...