Пример для распределенного обучения с использованием тензорного потока и повышенных деревьев

Question

Я прошел через многочисленные github, но я не могу найти пример, где распределенное обучение Tensorflow использовалось с оценщиком Boosted Trees Classifier. Все учебники для нейронных сетей.

Я немного адаптировал код расширенных деревьев для работы с распределенной стратегией, как показано ниже:

import numpy as np
import pandas as pd
from IPython.display import clear_output
from matplotlib import pyplot as plt
import tensorflow as tf
tf.random.set_seed(123)

dftrain = pd.read_csv('https://storage.googleapis.com/tf-datasets/titanic/train.csv')
dfeval = pd.read_csv('https://storage.googleapis.com/tf-datasets/titanic/eval.csv')
y_train = dftrain.pop('survived')
y_eval = dfeval.pop('survived')

fc = tf.feature_column
CATEGORICAL_COLUMNS = ['sex', 'n_siblings_spouses', 'parch', 'class', 'deck',
                       'embark_town', 'alone']
NUMERIC_COLUMNS = ['age', 'fare']

def one_hot_cat_column(feature_name, vocab):
  return tf.feature_column.indicator_column(
      tf.feature_column.categorical_column_with_vocabulary_list(feature_name,
                                                 vocab))
feature_columns = []
for feature_name in CATEGORICAL_COLUMNS:
  # Need to one-hot encode categorical features.
  vocabulary = dftrain[feature_name].unique()
  feature_columns.append(one_hot_cat_column(feature_name, vocabulary))

for feature_name in NUMERIC_COLUMNS:
  feature_columns.append(tf.feature_column.numeric_column(feature_name,
                                           dtype=tf.float32))

NUM_EXAMPLES = len(y_train)

def make_input_fn(X, y, n_epochs=None, shuffle=True):
  def input_fn():
    dataset = tf.data.Dataset.from_tensor_slices((dict(X), y))
    if shuffle:
      dataset = dataset.shuffle(NUM_EXAMPLES)
    # For training, cycle thru dataset as many times as need (n_epochs=None).
    dataset = dataset.repeat(n_epochs)
    # In memory training doesn't use batching.
    dataset = dataset.batch(NUM_EXAMPLES)
    return dataset
  return input_fn

# Training and evaluation input functions.
train_input_fn = make_input_fn(dftrain, y_train)
eval_input_fn = make_input_fn(dfeval, y_eval, shuffle=False, n_epochs=1)

n_batches = 10
mirrored_strategy = tf.distribute.MirroredStrategy()
config = tf.estimator.RunConfig(
    train_distribute=mirrored_strategy, eval_distribute=mirrored_strategy)
est = tf.estimator.BoostedTreesClassifier(feature_columns,
                                          n_batches_per_layer=n_batches,config=config)

  # The model will stop training once the specified number of trees is built, not
  # based on the number of steps.
est.train(train_input_fn, max_steps=100)

# Eval.
result = est.evaluate(eval_input_fn)
clear_output()
print(pd.Series(result))

Однако всякий раз, когда я запускаю этот код, я получаю сообщение об ошибке:

merge_call вызывается при определении нового графика или функции tf.function. Это часто может произойти, если функция fn, переданная в strategy.experimental_run(), украшена @tf.function (или содержит вложенный @tf.function), а fn содержит точку синхронизации, такую ​​как агрегирующие градиенты. Это поведение еще не поддерживается. Вместо этого, пожалуйста, оберните весь вызов strategy.experimental_run(fn) в @tf.function и избегайте вложенных tf.function s, которые потенциально могут пересекать границу синхронизации.

Итак, я был бы признателен, если бы я мог либо найдите способ отладки этой ошибки или найдите пример, который использует распределенное обучение с Boosted Trees.