The following sections provide more details about the dataset, running training and inference, and the training results.

Scripts and sample code

In the root directory, the most important files are:

train.py: Entry point for training data_utils.py: File containing the dataset implementation and preprocessing functions modeling.py: Definition of the model configuration.py: Contains configuration classes for various experiments test.py: Entry point testing trained model. Dockerfile: Container definition log_helper.py: Contains helper functions for setting up dllogger criterions.py: Definitions of loss functions

The scripts directory contains scripts for default use cases: run_electricity.sh: train default model on the electricity dataset run_traffic.sh: train default model on the traffic dataset

Command-line options

To view the full list of available options and their descriptions, use the -h or --help command-line option, for example: python train.py --help.

The following example output is printed when running the model:

usage: train.py [-h] --data_path DATA_PATH --dataset {electricity,traffic} [--epochs EPOCHS] [--sample_data SAMPLE_DATA SAMPLE_DATA] [--batch_size BATCH_SIZE] [--lr LR] [--seed SEED] [--use_amp] [--clip_grad CLIP_GRAD]
                [--early_stopping EARLY_STOPPING] [--results RESULTS] [--log_file LOG_FILE] [--distributed_world_size N] [--distributed_rank DISTRIBUTED_RANK] [--local_rank LOCAL_RANK] [--overwrite_config OVERWRITE_CONFIG]

optional arguments:
  -h, --help            show this help message and exit
  --data_path DATA_PATH
  --dataset {electricity,traffic}
  --epochs EPOCHS
  --sample_data SAMPLE_DATA SAMPLE_DATA
  --batch_size BATCH_SIZE
  --lr LR
  --seed SEED
  --use_amp             Enable automatic mixed precision
  --clip_grad CLIP_GRAD
  --early_stopping EARLY_STOPPING
                        Stop training if validation loss does not improve for more than this number of epochs.
  --results RESULTS
  --log_file LOG_FILE
  --distributed_world_size N
                        total number of GPUs across all nodes (default: all visible GPUs)
  --distributed_rank DISTRIBUTED_RANK
                        rank of the current worker
  --local_rank LOCAL_RANK
                        rank of the current worker
  --overwrite_config OVERWRITE_CONFIG
                        JSON string used to overload config

Getting the data

The TFT model was trained on the electricity and traffic benchmark datasets. This repository contains the get_data.sh download script, which for electricity and and traffic datasets will automatically download and preprocess the training, validation and test datasets, and produce files that contain scalers.

Dataset guidelines

The data_utils.py file contains all functions that are used to preprocess the data. Initially the data is loaded to a pandas.DataFrame and parsed to the common format which contains the features we will use for training. Then standardized data is cleaned, normalized, encoded and binarized. This step does the following: Drop all the columns that are not marked in the configuration file as used for training or preprocessing Flatten indices in case time series are indexed by more than one column Split the data into training, validation and test splits Filter out all the time series shorter than minimal example length Normalize columns marked as continuous in the configuration file Encode as integers columns marked as categorical Save the data in csv and binary formats

Multi-dataset

In order to use an alternate dataset, you have to write a function that parses your data to a common format. The format is as follows: There is at least one id column There is exactly one time column (that can also be used as a feature column) Each feature is in a separate column Each row represents a moment in time for only one time series Additionally, you must specify a configuration of the network, including a data description. Refer to the example in configuration.py file.

Training process

The train.py script is an entry point for a training procedure. Refined recipes can be found in the scripts directory. The model trains for at most --epochs epochs. If option --early_stopping N is set, then training will end if for N subsequent epochs validation loss hadn't improved. The details of the architecture and the dataset configuration are encapsulated by the --dataset option. This option chooses one of the configurations stored in the configuration.py file. You can enable mixed precision training by providing the --use_amp option. The training script supports multi-GPU training with the APEX package. To enable distributed training prepend training command with python -m torch.distributed.run --nproc_per_node=${NGPU}.

Example command:

python -m torch.distributed.run --nproc_per_node=8 train.py \
        --dataset electricity \
        --data_path /data/processed/electricity_bin \
        --batch_size=1024 \
        --sample 450000 50000 \
        --lr 1e-3 \
        --epochs 25 \
        --early_stopping 5 \
        --seed 1 \
        --use_amp \
        --results /results/TFT_electricity_bs8x1024_lr1e-3/seed_1

The model is trained by optimizing quantile loss . After training, the checkpoint with the least validation loss is evaluated on a test split with q-risk metric . Results are by default stored in the /results directory. This can be changed by providing the --results option. At the end of the training, the results directory will contain the trained checkpoint which had the lowest validation loss, dllogger logs (in dictionary per line format), and TensorBoard logs.

Inference process

Inference can be run by launching the inference.py script. The script requires a trained checkpoint to run. It is crucial to prepare the data in the same way as training data prior to running the inference. Example command:

python inference.py \
--checkpoint /results/checkpoint.pt \
--data /data/processed/electricity_bin/test.csv \
--tgt_scalers /data/processed/electricity_bin/tgt_scalers.bin \
--cat_encodings /data/processed/electricity_bin/cat_encodings.bin \
--batch_size 2048 \
--visualize \
--save_predictions \
--joint_visualization \
--results /results

In the default setting, it performs the evaluation of the model on a specified dataset and prints q-risk evaluated on this dataset. In order to save the predictions, use the --save_predictions option. Predictions will be stored in the directory specified by the --results option in the csv format. Option --joint_visualization allows us to plot graphs in TensorBoard format, allowing us to inspect the results and compare them to true values. Using --visualize, you can save plots for each example in a separate file.

Triton deployment

The NVIDIA Triton Inference Server provides a cloud inferencing solution optimized for NVIDIA GPUs. The server provides an inference service via an HTTP or GRPC endpoint, allowing remote clients to request inferencing for any model being managed by the server. More information on how to perform inference using NVIDIA Triton Inference Server can be found in triton/README.md.