Implementation:Recommenders team Recommenders NCF NNI Training

Knowledge Sources	Recommenders
Domains	Hyperparameter Tuning, Collaborative Filtering, NNI
Last Updated	2026-02-10 00:00 GMT

Overview

NNI trial script that trains and evaluates a Neural Collaborative Filtering (NCF) model with hyperparameters provided by the NNI automated tuning framework.

Description

This module serves as the NNI trial entry point for hyperparameter tuning of the NCF (NeuMF) recommendation model. When executed as a main script, it retrieves hyperparameters from NNI via nni.get_next_parameter(), merges them with command-line arguments parsed by get_params(), and supports Hyperband's STEPS parameter for dynamic epoch allocation.

The ncf_training function loads pickled train/validation DataFrames, instantiates an NCF model with NeuMF architecture (layer sizes [16, 8, 4]), trains the model, and evaluates it using both rating metrics (per-pair predictions on the validation set) and ranking metrics (all user-item predictions with seen items filtered via outer merge). Results are reported to NNI via nni.report_final_result() and persisted as JSON to the NNI_OUTPUT_DIR environment variable location, with the primary metric stored under the "default" key as required by NNI.

The helper function _update_metrics manages the metric dictionary, routing the primary metric to the "default" key and storing secondary metrics under their own names.

Usage

Use this module as the trial code file in an NNI experiment configuration for tuning NCF hyperparameters. It is designed to be launched by NNI's trial scheduler rather than invoked directly, though it can be tested standalone with appropriate command-line arguments and NNI environment setup.

Code Reference

Source Location

• Repository: Recommenders
• File: recommenders/tuning/nni/ncf_training.py
• Lines: 1-173

Signature

def _update_metrics(metrics_dict, metric, params, result)

def ncf_training(params)

def get_params()

def main(params)

Import

from recommenders.tuning.nni.ncf_training import ncf_training, get_params, main

I/O Contract

Inputs

ncf_training

Name	Type	Required	Description
params	dict	Yes	Dictionary containing all training configuration including: "datastore" (str), "train_datapath" (str), "validation_datapath" (str), "n_factors" (int), "n_epochs" (int), "learning_rate" (float), "verbose" (bool), "rating_metrics" (list), "ranking_metrics" (list), "k" (int), "primary_metric" (str)

get_params

Name	Type	Required	Description
--datastore	str	Yes	Path to the data directory
--train-datapath	str	Yes	Relative path to the pickled training data
--validation-datapath	str	Yes	Relative path to the pickled validation data
--rating-metrics	str (nargs)	No	List of rating metric names to evaluate
--ranking-metrics	str (nargs)	No	List of ranking metric names to evaluate
--k	int	No	Top-K value for ranking metrics
--epochs	int	No	Number of training epochs (default: 30)
--n_factors	int	No	Number of latent factors (default: 100)
--primary-metric	str	No	Primary metric name for NNI (default: "rmse")
--verbose	flag	No	Enable verbose output

Outputs

Name	Type	Description
return (ncf_training)	NCF model	The trained NCF model object
NNI report	dict	Metrics dictionary reported to NNI via nni.report_final_result()
metrics.json	file	JSON file with evaluation metrics written to NNI_OUTPUT_DIR

Usage Examples

Basic Usage

# This module is typically executed by NNI as a trial script.
# In an NNI config YAML:
# trial:
#   command: python ncf_training.py --datastore /data --train-datapath train.pkl --validation-datapath val.pkl --rating-metrics rmse mae --primary-metric rmse --epochs 50
#   codeDir: .

# Programmatic usage (for testing):
from recommenders.tuning.nni.ncf_training import ncf_training

params = {
    "datastore": "/path/to/data",
    "train_datapath": "train.pkl",
    "validation_datapath": "val.pkl",
    "n_factors": 64,
    "n_epochs": 20,
    "learning_rate": 0.001,
    "verbose": True,
    "rating_metrics": ["rmse", "mae"],
    "ranking_metrics": ["ndcg_at_k", "precision_at_k"],
    "k": 10,
    "primary_metric": "rmse",
}

model = ncf_training(params)