kumoai.trainer#

A Kumo Trainer supports training custom, highly performant graph neural network models atop a Graph and a TrainingTable, and prediction of these models with a PredictionTable. Models can be completely customized with detailed granularity with ModelPlan, although the default model plan suggested by predictive query is often suitable for great performance out-of-the-box.

Model Plan#

A ModelPlan defines the full parameter specification for training a Kumo model. It is broken down into multiple individual plans for different logical components of the training procedure: the ColumnProcessingPlan specifies any Encoder overrides for individual table columns, the ModelArchitecturePlan specifies graph neural network model parameters, the NeighborSamplingPlan specifies graph neural network subgraph sampling parameters, the OptimizationPlan specifies machine learning optimization parameters, and the TrainingJobPlan specifies training job-wide parameters for Kumo AutoML.

Note

After generating a default model plan with the suggest_model_plan() method, no further changes are necessary to train your first model using Kumo-inferred parameters. These options are provided in case you would like to further fine-tune the modeling plan.

ModelPlan

A complete definition of a Kumo model plan, encompassing a TrainingJobPlan, ColumnProcessingPlan, NeighborSamplingPlan, OptimizationPlan, and a ModelArchitecturePlan.

ColumnProcessingPlan

Configuration parameters that define how columns are encoded in the training and batch prediction pipelines.

ModelArchitecturePlan

Configuration parameters that define how the Kumo graph neural network is architected.

NeighborSamplingPlan

Configuration parameters that define how subgraphs are sampled in the training and batch prediction pipelines.

OptimizationPlan

Configuration parameters that define how columns are encoded in the training and batch prediction pipelines.

TrainingJobPlan

Configuration parameters that define the general execution of a Kumo AutoML search.

Training#

Training a model requires constructing a Trainer object atop a Graph and a TrainingTable. Fitting a model produces a TrainingJobResult or a TrainingJob that can be awaited at a later point in time.

Trainer

A trainer supports creating a Kumo machine learning model on a PredictiveQuery.

TrainingJob

Represents an in-progress training job.

TrainingJobResult

Represents a completed training job.

Batch Prediction#

Training a model requires constructing a Trainer object atop a Graph, PredictionTable, and a trained model (see the Trainer load method for more information). Predicting with a trained model produces a BatchPredictionJobResult or a BatchPredictionJob that can be awaited at a later point in time.

BatchPredictionJob

Represents an in-progress batch prediction job.

BatchPredictionJobResult

Represents a completed batch prediction job.