python - 在 Spark 中使用累积梯度进行分布式训练？

Question

我有一个在 SparkTF中Keras运行的自动编码器，目前该模型是在没有分布的驱动程序和完整数据集的一部分上本地训练的。

我想并行运行训练以加快过程并避免在整个数据集上出现 OOM。我的问题是，我只能在数据框上训练我的模型Pandas。因此，有时会在驱动程序上收集数据。

def keras_autoencoder(trainning_data, testing_data, out_model):
    # Prepare train and test data
    train_data = trainning_data.dataframe().where(F.col("msn") == F.lit("50038"))
    test_data = testing_data.dataframe().where(F.col("msn") == F.lit("50024"))

    # Create Processing Stage turning df into list of tensor
    column_transformer = Create_Tensor_EMA(outputCol="Input_Tensor", model_type="LSTM", column_order=column_order_EMA)
    vectorizer = Stage(column_transformer)

    training_df = vectorizer.transform(train_data)
    test_df = vectorizer.transform(test_data)
    Y = np.array(list(test_df.select('Input_Tensor').toPandas()['Input_Tensor']))
    X = np.array(list(training_df.select('Input_Tensor').toPandas()['Input_Tensor']))

    autoencoder = train_model(X, Y, EPOCHS, BATCH_SIZE, flight_len, param_len)
    sequential_stage = Stage(KerasModelWrapper(autoencoder), input_column_names='Input_Tensor', output_column_names='predictions')

    trainned_model = Model(vectorizer, sequential_stage)
    trainned_model.save(out_model)

我正在考虑使用rddand MapPartitions，因为我的数据已经按特定键进行了分区，因此可以将训练过程拆分到多个执行程序中。但我不希望每个执行者都在模型的单独实例上工作。

有没有办法使用gradient accumulation：

1. 让每个执行者在不更新权重的情况下训练一部分数据
2. 将生成的梯度发送回驱动程序以更新模型的权重
3. 将模型上的这个更新版本发回给执行器以处理下一批

标签： pythontensorflowkeraspysparkdeep-learning