Training Script Execution

Index SDK provides two training script execution modes:

  • (Recommended) trainCodeBook of the IVFSP algorithm
  • ivfsp_train_codebook.py The training script is stored in the tools/train folder in the installation directory. Note that the Python version is 3.9.11. To facilitate execution, ivfsp_train_codebook_example.sh is provided (stored in tools/train in the installation directory). You can modify parameter values in the script based on your requirements and run the script to generate a codebook file.

Command Reference

python3 ivfsp_train_codebook.py --dim <dim> --nonzero_num <nonzero_num> --nlist <nlist> --num_iter <num_iter> --device <device> --batch_size <batch_size> --code_num <code_num> --ratio <ratio> --learn_data_path <learn_data_path> --codebook_output_dir <codebook_output_dir> --train_model_dir <train_model_dir>

Parameter

<dim>: feature vector dimension The value must be consistent with that of <dim> generated by the training operator model file in IVFSP and must be greater than 0.

<nonzero_num>: number of non-zero dimensions after feature vector compression. The value must be consistent with that of <low_dim> generated by the training operator model file in IVFSP and must be greater than 0.

<nlist>: number of cluster centroids. The value must be consistent with that of <k> generated by the training operator model file in IVFSP and must be greater than 0.

<num_iter>: number of training iterations. The default value is 20. If the number of iterations is set to a large value, the training duration increases. The value must be greater than 0.

<device>: logical ID of a device. Training is performed on a specified device. The default value is 0.

<batch_size>: Training is performed based on its value. The value must be consistent with that of <batch_size> generated by the training operator model file in IVFSP and must be greater than 0. The default value is 32768.

<code_num>: maximum number of codebooks to be operated. The value must be a power of 2. The value must be consistent with that of <codebook_batch_size> generated by the training operator model file in IVFSP and must be greater than 0. The default value is 32768.

<ratio>: sampling rate of original samples used for training. The value range is (0, 1.0] and defaults to 1.0.

<learn_data_path>: path of the original feature file used for training. The bin and npy formats are supported. The bin storage mode is row-major first, and the data type is float32.

<codebook_output_dir>: directory to which the generated codebook file is exported. Ensure that this directory exists and the user who executes the program has the write permission on this directory. For security hardening, this directory level cannot contain soft links.

<train_model_dir>: directory where the IVFSP training operator model file is stored.

--help | -h: help information.

Instruction

  • After this command is executed, codebook_<dim>_<nonzero_num>_<nlist>.bin and codebook_<dim>_<nonzero_num>_<nlist>.npy are generated in the directory specified by <codebook_output_dir>. codebook_<dim>_<nonzero_num>_<nlist>.bin is the codebook file required by AscendIndexIVFSP.
  • When a codebook file exists, it will be overwritten. In this case, the program running user should be the owner of the file.
  • Before generating a training codebook, generate a training operator model file by referring to Generating the IVFSP Training Operator Model File.
  • The data size specified by learn_data_path must be greater than or equal to nonzero_num × nlist × sizeof(float32) bytes.