Question

我在 Spark 中运行的 TF 中有一个 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)

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

有没有办法使用 gradient accumulation 来：

让每个执行器对一部分数据进行训练而不更新权重
将生成的梯度发送回驱动程序以更新模型的权重
将模型上的此更新版本发送回执行程序以处理下一批

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

0 个答案: