NVIDIA TAO Toolkit v4.0
NVIDIA TAO Release tlt.40

Text Classification

Text Classification Model is a sequence classification model based on BERT-based encoders. It can be used for a variety of tasks like text classification, sentiment analysis, domain/intent detection for dialogue systems, etc. The model takes a text input and classifies it into predefined categories. Most of the BERT-based encoders supported by HuggingFace including BERT, RoBERTa, DistilBERT, XLNet, etc can be used with this model.

TAO Toolkit provides a sample notebook to outline the end-to-end workflow on how to train a Text classification model using TAO Toolkit and deploy it in Riva format on NGC resources.

Before proceeding, let’s download sample spec files that we would need for the rest of the subtasks.

Copy
Copied!
            

tao text_classification download_specs -r /results/text_classification/default_specs/ \ -o /specs/nlp/text_classification

where

  • -r : path to the logs that the command generates

  • -o : path where the spec files are downloaded

The text classification model uses NeMo toolkit and supports NeMo’s data format. It requires the data to be stored in TAB separated files (.tsv) with two columns of sentence and label. Each line of the data file contains text sequences, where words are separated with spaces and the label is separated with [TAB], i.e.:

Copy
Copied!
            

[WORD][SPACE][WORD][SPACE][WORD][TAB][LABEL]

Labels need to be integers starting from 0. Some examples taken from SST2 dataset, which is a two-class problem:

Copy
Copied!
            

saw how bad this movie was 0 lend some dignity to a dumb story 0 the greatest musicians 1

You may need three separate files for train, validation and test with this format.

If your dataset is stored in another format, you need to convert it to this format to use this model. There are some conversion scripts available for datasets SST2 and IMDB to convert them from their original format to NeMo’s format. To convert the original dataset to NeMo’s format, you can use dataset_convert utility as the following example:

Copy
Copied!
            

tao text_classification dataset_convert \ -e /specs/nlp/text_classification/dataset_convert.yaml \ source_data_dir=SOURCE_PATH \ target_data_dir=TARGET_PATH \ dataset_name=sst2

Arguments:

  • source_data_dir: directory of your dataset

  • target_data_dir: directory to save the converted dataset

  • datase_name: name of the dataset to convert (“imdb” and “sst2” are currently supported)

It converts the SST2 dataset stored as SOURCE_PATH to NeMo’s format and saves the new dataset at TARGET_PATH.

You may download the SST2 dataset from ‘https://dl.fbaipublicfiles.com/glue/data/SST-2.zip’ and extract it into the folder specified by SOURCE_PATH. After the conversion, the TARGET_PATH should contain the following files:

Copy
Copied!
            

. |--TARGET_PATH |-- train.tsv |-- dev.tsv

Example of config spec for training train.yaml file. You can change any of these parameters and pass them to the training command.

Copy
Copied!
            

trainer: max_epochs: 100 model: # Labels that will be used to "decode" predictions. class_labels: class_labels_file : null # optional to specify a file containing the list of the labels tokenizer: tokenizer_name: ${model.language_model.pretrained_model_name} # or sentencepiece vocab_file: null # path to vocab file tokenizer_model: null # only used if tokenizer is sentencepiece special_tokens: null language_model: pretrained_model_name: bert-base-uncased lm_checkpoint: null config_file: null # json file, precedence over config config: null classifier_head: # This comes directly from number of labels/target classes. num_output_layers: 2 fc_dropout: 0.1 training_ds: file_path: ??? batch_size: 64 shuffle: true num_samples: -1 # number of samples to be considered, -1 means all the dataset num_workers: 3 drop_last: false pin_memory: false validation_ds: file_path: ??? batch_size: 64 shuffle: false num_samples: -1 # number of samples to be considered, -1 means all the dataset num_workers: 3 drop_last: false pin_memory: false optim: name: adam lr: 2e-5 # optimizer arguments betas: [0.9, 0.999] weight_decay: 0.001 # scheduler setup sched: name: WarmupAnnealing # Scheduler params warmup_steps: null warmup_ratio: 0.1 last_epoch: -1 # pytorch lightning args monitor: val_loss reduce_on_plateau: false

Example of the command for training the model on four GPUs for 50 epochs:

Copy
Copied!
            

tao text_classification train -e /specs/nlp/text_classification/train.yaml \ training_ds.file_path=PATH_TO_TRAIN_FILE \ trainer.max_epochs=50 \ -g 4 \ -k $KEY

By default, the final model after training is done is saved in ‘trained-model.tlt’.

Required Arguments for Training

  • -e: The experiment specification file to set up training.

  • training_ds.file_path: Path to the training ‘.tsv’ file

  • -k: Encryption key

Optional Arguments

  • trainer.max_epochs: Training epochs number.

  • -g: Number of GPUs to use for training

  • Other arguments to override fields in the specification file.

The following table lists some of the parameters you may use in the config files and set them from command line when training a model:

Parameter

Data Type

Default

Description

model.class_labels.class_labels_file

string

null

Path to an optional file containing the labels; each line is the string label corresponding to a label

model.intent_loss_weight

float

0.6

Relation of intent to slot loss in total loss

model.tokenizer.tokenizer_name

string

Will be filled automatically based on model.language_model.pretrained_model_name

Tokenizer name

model.tokenizer.vocab_file

string

null

Path to tokenizer vocabulary

model.tokenizer.tokenizer_model

string

null

Path to tokenizer model (only for sentencepiece tokenizer)

model.tokenizer.special_tokens

string

null

Special tokens of the tokenizer if it exists

model.language_model.max_seq_length

integer

50

Maximal length of the input queries (in tokens)

model.language_model.pretrained_model_name

string

bert-base-uncased

Pre-trained language model name (choose from bert-base-cased, bert-base-uncased,
megatron_bert_345m_uncased, distilbert-base-uncased and biomegatron-bert-345m-uncased

model.language_model.lm_checkpoint

string

null

Path to the pre-trained language model checkpoint

model.language_model.config_file

string

null

Path to the pre-trained language model config file

model.language_model.config

dictionary

null

Config of the pre-trained language model

model.head.num_output_layers

integer

2

Number of fully connected layers of the Classifier on top of Bert model

model.head.fc_dropout

float

0.1

Dropout ratio of the fully connected layers

{training,validation,test}_ds.file_path

string

??

Path of the training ‘.tsv file

{training,validation,test}_ds.batch_size

integer

32

Data loader’s batch size

{training,validation,test}_ds.num_workers

integer

2

Number of worker threads for data loader

{training,validation,test}_ds.shuffle

boolean

true (training), false (test and validation)

Shuffles data for each epoch

{training,validation,test}_ds.drop_last

boolean

false

Specifies if last batch of data needs to get dropped if it is smaller than batch size

{training,validation,test}_ds.pin_memory

boolean

false

Enables pin_memory of PyTorch’s data loader to enhance speed

{training,validation,test}_ds.num_samples

integer

-1

Number of samples to be used from the dataset; -1 means all samples

optim.name

string

adam

Optimizer to use for training

optim.lr

float

2e-5

Learning rate to use for training

optim.weight_decay

float

0.01

Weight decay to use for training

optim.sched.name

string

WarmupAnnealing

Warmup schedule

optim.sched.warmup_ratio

float

0.1

Warmup ratio

Training Procedure

At the start of each training experiment, TAO Toolkit will print out a log of the experiment specification, including any parameters added or overridden via the command line. It will also show additional information, such as which GPUs are available and where logs will be saved. Then it shows some samples from the datasets with their corresponding inputs to the model.

Copy
Copied!
            

GPU available: True, used: True TPU available: None, using: 0 TPU cores LOCAL_RANK: 0 - CUDA_VISIBLE_DEVICES: [0,1,2,3] [NeMo W 2021-01-20 19:49:30 exp_manager:304] There was no checkpoint folder at checkpoint_dir :/home/user/tao-toolkit-pyt/nlp/text_classification/entrypoint/nemo_experiments/trained-model/2021-01-20_19-49-30/checkpoints. Training from scratch. [NeMo I 2021-01-20 19:49:30 exp_manager:194] Experiments will be logged at /home/user/tao-toolkit-pyt/nlp/text_classification/entrypoint/nemo_experiments/trained-model/2021-01-20_19-49-30

Then for each dataset provided, it shows some samples from the dataset with their corresponding inputs to the model. It also provides some stats on the lengths of sequences in the dataset.

Copy
Copied!
            

[NeMo I 2021-01-20 19:49:36 text_classification_dataset:120] Read 67350 examples from ../data/SST-2/train.tsv. [NeMo I 2021-01-20 19:49:37 text_classification_dataset:233] *** Example *** [NeMo I 2021-01-20 19:49:37 text_classification_dataset:234] example 0: ['girl-meets-girl', 'romantic', 'comedy'] [NeMo I 2021-01-20 19:49:37 text_classification_dataset:235] subtokens: [CLS] girl - meets - girl romantic comedy [SEP] [NeMo I 2021-01-20 19:49:37 text_classification_dataset:236] input_ids: 101 2611 1011 6010 1011 2611 6298 4038 102 [NeMo I 2021-01-20 19:49:37 text_classification_dataset:237] segment_ids: 0 0 0 0 0 0 0 0 0 [NeMo I 2021-01-20 19:49:37 text_classification_dataset:238] input_mask: 1 1 1 1 1 1 1 1 1 [NeMo I 2021-01-20 19:49:37 text_classification_dataset:239] label: 1

Before training starts, information on the optimizer and scheduler will be shown in the logs:

Copy
Copied!
            

[NeMo I 2021-01-20 19:50:19 modelPT:830] Optimizer config = Adam ( Parameter Group 0 amsgrad: False betas: [0.9, 0.999] eps: 1e-08 lr: 2e-05 weight_decay: 0.01 ) [NeMo I 2021-01-20 19:50:19 lr_scheduler:621] Scheduler "<nemo.core.optim.lr_scheduler.WarmupAnnealing object at 0x7fcd2232b160>" will be used during training (effective maximum steps = 1053) - Parameters : (warmup_steps: null warmup_ratio: 0.1 last_epoch: -1 max_steps: 1053 )

You should next see a full printout of the number of parameters in each module and submodule, as well as the total number of trainable and non-trainable parameters in the model. For example, this model has 100M parameters in total:

Copy
Copied!
            

| Name | Type | Params -------------------------------------------------------------------------------------------------- 0 | bert_model | BertEncoder | 109 M 1 | bert_model.embeddings | BertEmbeddings | 23.8 M 2 | bert_model.embeddings.word_embeddings | Embedding | 23.4 M 3 | bert_model.embeddings.position_embeddings | Embedding | 393 K 4 | bert_model.embeddings.token_type_embeddings | Embedding | 1.5 K 5 | bert_model.embeddings.LayerNorm | LayerNorm | 1.5 K 6 | bert_model.embeddings.dropout | Dropout | 0 7 | bert_model.encoder | BertEncoder | 85.1 M 8 | bert_model.encoder.layer | ModuleList | 85.1 M 9 | bert_model.encoder.layer.0 | BertLayer | 7.1 M 10 | bert_model.encoder.layer.0.attention | BertAttention | 2.4 M 11 | bert_model.encoder.layer.0.attention.self | BertSelfAttention | 1.8 M 12 | bert_model.encoder.layer.0.attention.self.query | Linear | 590 K ... 212 | bert_model.encoder.layer.11.output.dropout | Dropout | 0 213 | bert_model.pooler | BertPooler | 590 K 214 | bert_model.pooler.dense | Linear | 590 K 215 | bert_model.pooler.activation | Tanh | 0 216 | classifier | SequenceClassifier | 592 K 217 | classifier.dropout | Dropout | 0 218 | classifier.mlp | MultiLayerPerceptron | 592 K 219 | classifier.mlp.layer0 | Linear | 590 K 220 | classifier.mlp.layer2 | Linear | 1.5 K 221 | loss | CrossEntropyLoss | 0 222 | classification_report | ClassificationReport | 0 -------------------------------------------------------------------------------------------------- 110 M Trainable params 0 Non-trainable params 110 M Total params

As the model starts training, you should see a progress bar per epoch.

Copy
Copied!
            

Epoch 0: 100%|████████████████████████████| 1067/1067 [03:10<00:00, 5.60it/s, loss=0.252, val_loss=0.258, Epoch 0, global step 1052: val_loss reached 0.25792 (best 0.25792), saving model to "/home/user/tao-toolkit-pyt/nlp/text_classification/entrypoint/nemo_experiments/trained-model/2021-01-20_20-19-44/checkpoints/trained-model---val_loss=0.26-epoch=0.ckpt" as top 3 Epoch 1: 100%|████████████████████████████| 1067/1067 [03:10<00:00, 5.60it/s, loss=0.187, val_loss=0.245, Epoch 1, global step 2105: val_loss reached 0.24499 (best 0.24499), saving model to "/home/user/tao-toolkit-pyt/nlp/text_classification/entrypoint/nemo_experiments/trained-model/2021-01-20_20-19-44/checkpoints/trained-model---val_loss=0.24-epoch=1.ckpt" as top 3 Epoch 2: 100%|████████████████████████████| 1067/1067 [03:09<00:00, 5.62it/s, loss=0.158, val_loss=0.235, Epoch 2, global step 3158: val_loss reached 0.23505 (best 0.23505), saving model to "/home/user/tao-toolkit-pyt/nlp/text_classification/entrypoint/nemo_experiments/trained-model/2021-01-20_20-19-44/checkpoints/trained-model---val_loss=0.24-epoch=2.ckpt" as top 3 ...

After each epoch, you should see a summary table of metrics on the validation set.

Copy
Copied!
            

Validating: 100%|████████████████████████████| 14/14 [00:00<00:00, 13.94it/s] [NeMo I 2021-01-20 19:53:32 text_classification_model:173] val_report: label precision recall f1 support label_id: 0 91.97 88.32 90.11 428 label_id: 1 89.15 92.57 90.83 444 ------------------- micro avg 90.48 90.48 90.48 872 macro avg 90.56 90.44 90.47 872 weighted avg 90.54 90.48 90.47 872

At the end of training, TAO Toolkit will save the last checkpoint at the path specified by the experiment spec file before finishing.

Copy
Copied!
            

Saving latest checkpoint... [NeMo I 2021-01-20 21:09:39 train:124] Experiment logs saved to '/home/user/tao-toolkit-pyt/nlp/text_classification/entrypoint/nemo_experiments/trained-model/2021-01-20_21-06-17' [NeMo I 2021-01-20 21:09:39 train:127] Trained model saved to '/home/user/tao-toolkit-pyt/nlp/text_classification/entrypoint/nemo_experiments/trained-model/2021-01-20_21-06-17/checkpoints/trained-model.tlt'

The output logs for the evaluation and fine-tuning look similar.

When you want to train this model on other data or with different batch sizes, you may need to tune at least the configs of your optimizer and scheduler like the learning rate and weight decay. Higher effective batch sizes need larger learning rate. Effective batch size is the total number of your samples per each update step. For example, when your batch size per GPU is set to 64, and you use four GPUs with accumulate_grad_batches of two, then your effective batch size would be 512=64*4*2. You may use other BERT-like models or models with different sizes based on your performance requirements.

There are scenarios where users are required to re-train or fine-tune a pretrained TAO model like trained-model.tlt on a new dataset. TAO toolkit provides a separate tool called fine-tune to enable this.

Example of spec file to be used for fine-tuning of a model:

Copy
Copied!
            

trainer: max_epochs: 100 data_dir: ??? # Fine-tuning settings: training dataset. finetuning_ds: file_path: ??? batch_size: 64 shuffle: false num_samples: -1 # number of samples to be considered, -1 means all the dataset num_workers: 3 drop_last: false pin_memory: false # Fine-tuning settings: validation dataset. validation_ds: file_path: ??? batch_size: 64 shuffle: false num_samples: -1 # number of samples to be considered, -1 means all the dataset num_workers: 3 drop_last: false pin_memory: false # Fine-tuning settings: different optimizer. optim: name: adam lr: 2e-5 betas: [0.9, 0.9998] weight_decay: 0.001

Use the following command to fine-tune a pre-trained model on a training file specified by ‘finetuning_ds.file_path’:

Copy
Copied!
            

tao text_classification finetune [-h] -e /specs/nlp/text_classification/finetune.yaml \ -r PATH_TO_RESULT_FOLDER \ -m PATH_OF_PRETRAINED_TLT_MODEL \ -g 1 \ finetuning_ds.file_path=PATH_TO_TRAIN_FILE \ trainer.max_epochs=3 \ -k $KEY

Required Arguments for Fine-tuning

  • -e: The experiment specification file to set up fine-tuning

  • -r: Path to the directory to store the results of the fine-tuning.

  • -m: Path to the pre-trained model to use for fine-tuning.

  • training_ds.file_path: Path to the training ‘.tsv’ file

  • -k: Encryption key

Optional Arguments

  • -h, --help: Show this help message and exit

  • -g: The number of GPUs to be used in evaluation in a multi-GPU scenario (default: 1).

  • Other arguments to override fields in the specification file.

The evaluation tool enables the user to evaluate a saved model in TAO format on a dataset.

Spec example to evaluate the pre-trained model on test data:

Copy
Copied!
            

restore_from: trained-model.tlt test_ds: file_path: PATH_TO_TEST_FILE num_workers: 2 batch_size: 32 shuffle: false num_samples: -1

Use the following command to evaluate the model:

Copy
Copied!
            

tao text_classification evaluate \ -e /specs/nlp/text_classification/evaluate.yaml \ test_ds.file_path=PATH_TO_TEST_FILE \

Required Arguments for Evaluation

  • -e: The experiment specification file to set up evaluation.

  • test_ds.file_path: Path to the test ‘.tsv’ file

The output should be similar to the training process and the metrics table is shown at the end:

Copy
Copied!
            

Validating: 100%|████████████████████████████| 14/14 [00:00<00:00, 13.94it/s] [NeMo I 2021-01-20 19:53:32 text_classification_model:173] val_report: label precision recall f1 support label_id: 0 91.97 88.32 90.11 428 label_id: 1 89.15 92.57 90.83 444 ------------------- micro avg 90.48 90.48 90.48 872 macro avg 90.56 90.44 90.47 872 weighted avg 90.54 90.48 90.47 872

This table contains the metrics for each class separately, like precision, recall, F1, and support. It also shows Micro Average, Macro Average, and Weighted Average, which may show the overall performance of the model on all classes.

Inference tool would take some inputs in text format and produces the predictions of a saved model for them. To run inference on the model, specify the list of examples in the spec file “infer.yaml”, for example:

Copy
Copied!
            

input_batch: - "by the end of no such thing the audience , like beatrice , has a watchful affection for the monster ." - "director rob marshall went out gunning to make a great one ." - "uneasy mishmash of styles and genres ." - "I love exotic science fiction / fantasy movies but this one was very unpleasant to watch . Suggestions and images of child abuse , mutilated bodies (live or dead) , other gruesome scenes , plot holes , boring acting made this a regretable experience , The basic idea of entering another person's mind is not even new to the movies or TV (An Outer Limits episode was better at exploring this idea) . i gave it 4 / 10 since some special effects were nice ."

The list of inputs specified by ‘input_batch’ would be passed through the model to get the label predictions.

To run the inference on a trained model ‘trained-model.tlt’:

Copy
Copied!
            

tao text_classification infer \ -e /specs/nlp/text_classification/infer.yaml \ -m trained-model.tlt \

Required Arguments for Inference

  • -e: The experiment specification file to set up inference. This requires the input_batch with the list of examples to run inference on.

  • -m: Path to the pre-trained model checkpoint from which to infer. Should be a .tlt file.

The output should be like this:

Copy
Copied!
            

[NeMo I 2021-01-27 16:58:32 infer:68] Spec file: restore_from: trained-model.tlt exp_manager: task_name: infer explicit_log_dir: ./ input_batch: - by the end of no such thing the audience , like beatrice , has a watchful affection for the monster . - director rob marshall went out gunning to make a great one . - uneasy mishmash of styles and genres . - I love exotic science fiction / fantasy movies but this one was very unpleasant to watch . Suggestions and images of child abuse , mutilated bodies (live or dead) , other gruesome scenes , plot holes , boring acting made this a regretable experience , The basic idea of entering another person's mind is not even new to the movies or TV (An Outer Limits episode was better at exploring this idea) . i gave it 4 / 10 since some special effects were nice . encryption_key: null ... [NeMo I 2021-01-27 16:58:50 infer:101] Query: by the end of no such thing the audience , like beatrice , has a watchful affection for the monster . [NeMo I 2021-01-27 16:58:50 infer:102] Predicted label: positive [NeMo I 2021-01-27 16:58:50 infer:101] Query: director rob marshall went out gunning to make a great one . [NeMo I 2021-01-27 16:58:50 infer:102] Predicted label: positive [NeMo I 2021-01-27 16:58:50 infer:101] Query: uneasy mishmash of styles and genres . [NeMo I 2021-01-27 16:58:50 infer:102] Predicted label: negative [NeMo I 2021-01-27 16:58:50 infer:101] Query: I love exotic science fiction / fantasy movies but this one was very unpleasant to watch . Suggestions and images of child abuse , mutilated bodies (live or dead) , other gruesome scenes , plot holes , boring acting made this a regretable experience , The basic idea of entering another person's mind is not even new to the movies or TV (An Outer Limits episode was better at exploring this idea) . i gave it 4 / 10 since some special effects were nice . [NeMo I 2021-01-27 16:58:50 infer:102] Predicted label: negative

Each query would be printed out along with its predicted label.

You may use the export toolkit to convert a pre-trained saved TAO model into Riva format. This format would enable faster inference. An example of the spec file for model export:

Copy
Copied!
            

# Name of the .tlt EFF archive to be loaded/model to be exported. restore_from: trained-model.tlt # Set export format to RIVA export_format: RIVA # Output EFF archive containing Riva file. export_to: exported-model.riva

Parameter

Data Type

Default

Description

restore_from

string

trained-model.tlt

Path to the pre-trained model

export_format

string

ONNX

Export format, choose from: ONNX or RIVA

export_to

string

exported-model.eonnx

Path to the exported model

To export a pre-trained model to RIVA’s format, run:

Copy
Copied!
            

tao text_classification export \ -e /specs/nlp/text_classification/export.yaml \ -m finetuned-model.tlt \ -k $KEY \ export_format=RIVA \ export_to=exported-model.riva

Required Arguments for Export

  • -e: The experiment specification file to set up inference. This requires the input_batch with the list of examples to run inference on.

  • -m: Path to the pre-trained model checkpoint from which to infer. Should be a .tlt file.

  • -k: Encryption key

The output should be something like this:

Copy
Copied!
            

[NeMo I 2021-01-27 17:04:05 export:37] Spec file: restore_from: ./trained-model.tlt export_to: exported-model.riva export_format: RIVA exp_manager: task_name: export explicit_log_dir: ./ encryption_key: null ... [NeMo W 2021-02-01 16:00:22 exp_manager:27] Exp_manager is logging to `./``, but it already exists. [NeMo W 2021-02-01 16:00:28 modelPT:193] Using /tmp/tmpmke24h_1/tokenizer.vocab_file instead of tokenizer.vocab_file. [NeMo W 2021-02-01 16:00:35 modelPT:193] Using /tmp/tmpmke24h_1/label_ids.csv instead of ../data/SST-2/label_ids.csv. [NeMo I 2021-02-01 16:00:37 export:52] Model restored from '/home/user/tao-toolkit-pyt/nlp/text_classification/entrypoint/nemo_experiments/trained-model/2021-01-27_16-53-38/checkpoints/trained-model.tlt' [NeMo I 2021-02-01 16:01:08 export:66] Experiment logs saved to '.' [NeMo I 2021-02-01 16:01:08 export:67] Exported model to './exported-model.riva'


© Copyright 2022, NVIDIA.. Last updated on Mar 23, 2023.