Class KerasEstimator

horovod/spark/keras/estimator.py:98–385

Spark Estimator for fitting Keras models to a DataFrame. Supports standalone `keras` and `tf.keras`, and TensorFlow 1.X and 2.X. Args: num_proc: Number of Horovod processes. Defaults to `spark.default.parallelism`. data_module: (Optional) DataModule class used for training

class KerasEstimator(HorovodEstimator, KerasEstimatorParamsReadable,
                     KerasEstimatorParamsWritable):
    """Spark Estimator for fitting Keras models to a DataFrame.

    Supports standalone `keras` and `tf.keras`, and TensorFlow 1.X and 2.X.

    Args:
        num_proc: Number of Horovod processes.  Defaults to `spark.default.parallelism`.
        data_module: (Optional) DataModule class used for training and validation, if not set, defaults to the PetastormDataModule.
        model: Keras model to train.
        backend: Optional Backend object for running distributed training function. Defaults to SparkBackend with
                 `num_proc` worker processes. Cannot be specified if `num_proc` is also provided.
        store: Store object that abstracts reading and writing of intermediate data and run results.
        custom_objects: Optional dictionary mapping names (strings) to custom classes or functions to be considered
                        during serialization/deserialization.
        optimizer: Keras optimizer to be converted into a `hvd.DistributedOptimizer` for training.
        loss: Keras loss or list of losses.
        loss_weights: Optional list of float weight values to assign each loss.
        sample_weight_col: Optional column indicating the weight of each sample.
        gradient_compression: Gradient compression used by `hvd.DistributedOptimizer`.
        metrics: Optional metrics to record.
        feature_cols: Column names used as feature inputs to the model. Must be a list with each feature
                      mapping to a sequential argument in the model's forward() function.
        label_cols: Column names used as labels.  Must be a list with one label for each output of the model.
        validation: Optional validation column name (string) where every row in the column is either 1/True or 0/False,
                    or validation split (float) giving percent of data to be randomly selected for validation.
        callbacks: Keras callbacks.
        batch_size: Number of rows from the DataFrame per batch.
        val_batch_size: Number of rows from the DataFrame per batch for validation, if not set, will use batch_size.
        epochs: Number of epochs to train.
        verbose: Verbosity level [0, 2] (default: 1).
        random_seed: Optional random seed to use for Tensorflow. Default: None.
        shuffle_buffer_size: (Deprecated) Optional size of in-memory shuffle buffer in rows (on training data).
                             Allocating a larger buffer size increases randomness of shuffling at
                             the cost of more host memory. Defaults to estimating with an assumption
                             of 4GB of memory per host. Set shuffle_buffer_size=0 would turn off shuffle.
        shuffle: (Optional) Whether to shuffle training samples or not. Defaults to True.
        partitions_per_process: Number of Parquet partitions to assign per worker process from `num_proc` (default: 10).
        run_id: Optional unique ID for this run for organization in the Store. Will be automatically assigned if not
                provided.
        train_steps_per_epoch: Number of steps to train each epoch. Useful for testing that model trains successfully.
                               Defaults to training the entire dataset each epoch.
        validation_steps_per_epoch: Number of validation steps to perform each epoch.
        transformation_fn: Optional function that takes a row as its parameter
                           and returns a modified row that is then fed into the
                           train or validation step. This transformation is
                           applied after batching. See Petastorm [TransformSpec](https://github.com/uber/petastorm/blob/master/petastorm/transform.py)
                           for more details. Note that this fucntion constructs
                           another function which should perform the
                           transformation.
        train_reader_num_workers: This parameter specifies the number of parallel processes that
                               read the training data from data store and apply data
                               transformations to it. Increasing this number
                               will generally increase the reading rate but will also
                               increase the memory footprint. More processes are
                               particularly useful if the bandwidth to the data store is not
                               high enough, or users need to apply transformation such as
                               decompression or data augmentation on raw data.