ReIdentificationNet
ReIdentificationNet takes cropped images of a person from different perspectives as network input and outputs the embedding features for that person. The embeddings are used to perform similarity matching to re-identify the same person. The model supported in the current version is based on ResNet, which is the most commonly used baseline for re-identification due to its high accuracy.
The ReIdentificationNet apps in TAO Toolkit expect data in Market-1501 format for training and evaluation.
See the Data Annotation Format page for more information about the Market-1501 data format.
The spec file for ReIdentificationNet includes model
, dataset
,
re_ranking
, and train
parameters. Here is an example spec
for training a ResNet model on Market-1501 that contains 751 identities in the training set.
results_dir: "/path/to/experiment_results"
encryption_key: nvidia_tao
model:
backbone: resnet_50
last_stride: 1
pretrain_choice: imagenet
pretrained_model_path: "/path/to/pretrained_model.pth"
input_channels: 3
input_width: 128
input_height: 256
neck: bnneck
feat_dim: 256
neck_feat: after
metric_loss_type: triplet
with_center_loss: False
with_flip_feature: False
label_smooth: True
dataset:
train_dataset_dir: "/path/to/train_dataset_dir"
test_dataset_dir: "/path/to/test_dataset_dir"
query_dataset_dir: "/path/to/query_dataset_dir"
num_classes: 751
batch_size: 64
val_batch_size: 128
num_workers: 1
pixel_mean: [0.485, 0.456, 0.406]
pixel_std: [0.226, 0.226, 0.226]
padding: 10
prob: 0.5
re_prob: 0.5
sampler: softmax_triplet
num_instances: 4
re_ranking:
re_ranking: True
k1: 20
k2: 6
lambda_value: 0.3
train:
results_dir: "${results_dir}/train"
optim:
name: Adam
lr_monitor: val_loss
steps: [40, 70]
gamma: 0.1
bias_lr_factor: 1
weight_decay: 0.0005
weight_decay_bias: 0.0005
warmup_factor: 0.01
warmup_iters: 10
warmup_method: linear
base_lr: 0.00035
momentum: 0.9
center_loss_weight: 0.0005
center_lr: 0.5
triplet_loss_margin: 0.3
num_epochs: 120
checkpoint_interval: 10
Parameter | Data Type | Default | Description |
model |
dict config | – | The configuration for the model architecture |
train |
dict config | – | The configuration for the training process |
dataset |
dict config | – | The configuration for the dataset |
re_ranking |
dict config | – | The configuration for the re-ranking module |
model
The model
parameter provides options to change the ReIdentificationNet architecture.
model:
backbone: resnet_50
last_stride: 1
pretrain_choice: imagenet
pretrained_model_path: "/path/to/pretrained_model.pth"
input_channels: 3
input_width: 128
input_height: 256
neck: bnneck
feat_dim: 256
neck_feat: after
metric_loss_type: triplet
with_center_loss: False
with_flip_feature: False
label_smooth: True
Parameter | Datatype | Default | Description | Supported Values |
backbone |
string | resnet_50 | The type of model, which can be resnet_50 or a Swin-based architecture (refer to ReIdentificationNet Transformer for more details) | “resnet_50”, “swin_base_patch4_window7_224”, “swin_small_patch4_window7_224, “swin_tiny_patch4_window7_224” |
last_stride |
unsigned int | 1 | The number of strides during convolution | >0 |
pretrain_choice |
string | imagenet | The pre-trained network | imagenet/self/”” |
pretrained_model_path |
string | The path to the pre-trained model | ||
input_channels |
unsigned int | 3 | The number of input channels | >0 |
input_width |
int | 128 | The width of the input images | >0 |
input_height |
int | 256 | The height of the input images | >0 |
neck |
string | bnneck | Specifies whether to train with BNNeck | bnneck/”” |
feat_dim |
unsigned int | 256 | The output size of the feature embeddings | >0 |
neck_feat |
string | after | Specifies which feature of BNNeck to use for testing | before/after |
metric_loss_type |
string | triplet | The type of metric loss | triplet/center/triplet_center |
with_center_loss |
bool | False | Specifies whether to enable center loss | True/False |
with_flip_feature |
bool | False | Specifies whether to enable image flipping | True/False |
label_smooth |
bool | True | Specifies whether to enable label smoothing | True/False |
dataset
The dataset
parameter defines the dataset source, training batch size, and augmentation.
dataset:
train_dataset_dir: "/path/to/train_dataset_dir"
test_dataset_dir: "/path/to/test_dataset_dir"
query_dataset_dir: "/path/to/query_dataset_dir"
num_classes: 751
batch_size: 64
val_batch_size: 128
num_workers: 1
pixel_mean: [0.485, 0.456, 0.406]
pixel_std: [0.226, 0.226, 0.226]
padding: 10
prob: 0.5
re_prob: 0.5
sampler: softmax_triplet
num_instances: 4
Parameter | Datatype | Default | Description | Supported Values |
train_dataset_dir |
string | The path to the train images | ||
test_dataset_dir |
string | The path to the test images | ||
query_dataset_dir |
string | The path to the query images | ||
num_classes |
unsigned int | 751 | The number of unique person IDs | >0 |
batch_size |
unsigned int | 64 | The batch size for training | >0 |
val_batch_size |
unsigned int | 128 | The batch size for validation | >0 |
num_workers |
unsigned int | 1 | The number of parallel workers processing data | >0 |
pixel_mean |
float list | [0.485, 0.456, 0.406] | The pixel mean for image normalization | float list |
pixel_std |
float list | [0.226, 0.226, 0.226] | The pixel standard deviation for image normalization | float list |
padding |
unsigned int | 10 | The pixel padding size around images for image augmentation | >=1 |
prob |
float | 0.5 | The random horizontal flipping probability for image augmentation | >0 |
re_prob |
float | 0.5 | The random erasing probability for image augmentation | >0 |
sampler |
string | softmax_triplet | The type of sampler for data loading | softmax/triplet/softmax_triplet |
num_instances |
unsigned int | 4 | The number of image instances of the same person in a batch | >0 |
re_ranking
The re_ranking
parameter defines the settings for the re-ranking module.
re_ranking:
re_ranking: True
k1: 20
k2: 6
lambda_value: 0.3
Parameter | Datatype | Default | Description | Supported Values |
re_ranking |
bool | True | A flag that enables the re-ranking module | True/False |
k1 |
unsigned int | 20 | The k used for k-reciprocal nearest neighbors | >0 |
k2 |
unsigned int | 6 | The k used for local query expansion | >0 |
lambda_value |
float | 0.3 | The weight of original distance in the combination with Jaccard distance | >0.0 |
train
The train
parameter defines the hyperparameters of the training process.
train:
optim:
name: Adam
lr_monitor: val_loss
steps: [40, 70]
gamma: 0.1
bias_lr_factor: 1
weight_decay: 0.0005
weight_decay_bias: 0.0005
warmup_factor: 0.01
warmup_iters: 10
warmup_method: linear
base_lr: 0.00035
momentum: 0.9
center_loss_weight: 0.0005
center_lr: 0.5
triplet_loss_margin: 0.3
num_epochs: 120
checkpoint_interval: 10
Parameter | Datatype | Default | Description | Supported Values |
optim |
dict config | The configuration for the SGD optimizer, including the learning rate, learning scheduler, weight decay, etc. | ||
num_epochs |
unsigned int | 120 | The total number of epochs to run the experiment | >0 |
checkpoint_interval |
unsigned int | 10 | The interval at which the checkpoints are saved | >0 |
clip_grad_norm |
float | 0.0 | The amount to clip the gradient by the L2 norm. A value of 0.0 specifies no clipping. | >=0 |
optim
The optim
parameter defines the config for the SGD optimizer in training, including the
learning rate, learning scheduler, and weight decay.
optim:
name: Adam
lr_monitor: val_loss
lr_steps: [40, 70]
gamma: 0.1
bias_lr_factor: 1
weight_decay: 0.0005
weight_decay_bias: 0.0005
warmup_factor: 0.01
warmup_iters: 10
warmup_method: linear
base_lr: 0.00035
momentum: 0.9
center_loss_weight: 0.0005
center_lr: 0.5
triplet_loss_margin: 0.3
Parameter |
Datatype |
Default |
Description |
Supported Values |
---|---|---|---|---|
name |
string | Adam | The name of the optimizer | Adam/SGD/Adamax/… |
lr_monitor |
string | val_loss | The monitor value for the AutoReduce scheduler | val_loss/train_loss |
lr_steps |
int list | [40, 70] | The steps to decrease the learning rate for the MultiStep scheduler |
int list |
gamma |
float | 0.1 | The decay rate for the WarmupMultiStepLR | >0.0 |
bias_lr_factor |
float | 1 | The bias learning rate factor for the WarmupMultiStepLR | >=1 |
weight_decay |
float | 0.0005 | The weight decay coefficient for the optimizer | >0.0 |
weight_decay_bias |
float | 0.0005 | The weight decay bias for the optimizer | >0.0 |
warmup_factor |
float | 0.01 | The warmup factor for the WarmupMultiStepLR scheduler | >0.0 |
warmup_iters |
unsigned int | 10 | The number of warmup iterations for the WarmupMultiStepLR scheduler | >0 |
warmup_method |
string | linear | The warmup method for the optimizer | linear/cosine |
base_lr |
float | 0.00035 | The initial learning rate for the training | >0.0 |
momentum |
float | 0.9 | The momentum for the WarmupMultiStepLR optimizer | >0.0 |
center_loss_weight |
float | 0.0005 | The balanced weight of center loss | >0.0 |
center_lr |
float | 0.5 | The learning rate of SGD to learn the centers of center loss | >0.0 |
triplet_loss_margin |
float | 0.3 | The margin value for triplet loss | >0.0 |
Use the following command to run ReIdentificationNet training:
tao model re_identification train -e <experiment_spec_file>
-r <results_dir>
-k <key>
[train.gpu_ids=<gpu id list>]
Required Arguments
-e, --experiment_spec_file
: The path to the experiment spec file.-r, --results_dir
: The path to a folder where the experiment outputs should be written.-k, --key
: The user-specific encoding key to save or load a.tlt
model.
Optional Arguments
train.gpu_ids
: The GPU indices list for training. If you set more than one GPU ID, multi-GPU training will be triggered automatically.
Here’s an example of using the ReIdentificationNet training command:
tao model re_identification train -e $DEFAULT_SPEC -r $RESULTS_DIR -k $KEY
The evaluation metric of ReIdentificationNet is the mean average precision and ranked accuracy.
The plots of sampled matches and the cumulative matching characteristic (CMC) curve can be obtained using
the evaluate.output_sampled_matches_plot
and evaluate.output_cmc_curve_plot
parameters,
respectively.
Use the following command to run ReIdentificationNet evaluation:
tao model re_identification evaluate -e <experiment_spec_file>
-r <results_dir>
-k <key>
evaluate.checkpoint=<model to be evaluated>
evaluate.output_sampled_matches_plot=<path to the output sampled matches plot>
evaluate.output_cmc_curve_plot=<path to the output CMC curve plot>
evaluate.test_dataset=<path to test data>
evaluate.query_dataset=<path to query data>
[evaluate.gpu_id=<gpu index>]
Required Arguments
-e, --experiment_spec_file
: The experiment spec file to set up the evaluation experiment-r, --results_dir
: The path to a folder where the experiment outputs should be written-k, --key
: The encoding key for the.tlt
modelevaluate.checkpoint
: The.tlt
modelevaluate.output_sampled_matches_plot
: The path to the plotted file of sampled matchesevaluate.output_cmc_curve_plot
: The path to the plotted file of the CMC curveevaluate.test_dataset
: The path to the test dataevaluate.query_dataset
: The path to the query data
Optional Argument
evaluate.gpu_id
: The GPU index used to run the evaluation. You can specify the GPU index used to run evaluation when the machine has multiple GPUs installed. Note that evaluation can only run on a single GPU.
Here’s an example of using the ReIdentificationNet evaluation command:
tao model re_identification evaluate -e $DEFAULT_SPEC -r $RESULTS_DIR -k $KEY evaluate.checkpoint=$TRAINED_TLT_MODEL evaluate.output_sampled_matches_plot=$OUTPUT_SAMPLED_MATCHED_PLOT evaluate.output_cmc_curve_plot=$OUTPUT_CMC_CURVE_PLOT evaluate.test_dataset=$TEST_DATA evaluate.query_dataset=$QUERY_DATA
Use the following command to run inference on ReIdentificationNet with the .tlt
model.
tao model re_identification inference -e <experiment_spec>
-r <results_dir>
-k <key>
inference.checkpoint=<inference model>
inference.output_file=<path to output file>
inference.test_dataset=<path to gallery data>
inference.query_dataset=<path to query data>
[inference.gpu_id=<gpu index>]
The output will be a JSON file that contains the feature embeddings of all the test and query data.
Required Arguments
-e, --experiment_spec
: The experiment spec file to set up inference-r, --results_dir
: The path to a folder where the experiment outputs should be written-k, --key
: The encoding key for the.tlt
modelinference.checkpoint
: The.tlt
model to perform inference withinference.output_file
: The path to the output JSON fileinference.test_dataset
: The path to the test datainference.query_dataset
: The path to the query data
Optional Argument
inference.gpu_id
: The index of the GPU that will be used to run inference. You can specify this value when the machine has multiple GPUs installed. Note that inference can only run on a single GPU.
Here’s an example of using the ReIdentificationNet inference command:
tao model re_identification inference -e $DEFAULT_SPEC -r $RESULTS_DIR -k $KEY inference.checkpoint=$TRAINED_TLT_MODEL inference.output_file=$OUTPUT_FILE inference.test_dataset=$TEST_DATA inference.query_dataset=$QUERY_DATA
The expected output would be as follows:
[
{
"img_path": "/path/to/img1.jpg",
"embedding": [-0.30, 0.12, 0.13,...]
},
{
"img_path": "/path/to/img2.jpg",
"embedding": [-0.10, -0.06, -1.85,...]
},
...
{
"img_path": "/path/to/imgN.jpg",
"embedding": [1.41, 0.63, -0.15,...]
}
]
Use the following command to export ReIdentificationNet to .onnx
format for deployment:
tao model re_identification export -e <experiment_spec>
-r <results_dir>
-k <key>
export.checkpoint=<tlt checkpoint to be exported>
[export.onnx_file=<path to exported file>]
[export.gpu_id=<gpu index>]
Required Arguments
-e, --experiment_spec
: The experiment spec file to set up export.-r, --results_dir
: The path to a folder where the experiment outputs should be written.-k, --key
: The encoding key for the.tlt
model.export.checkpoint
: The.tlt
model to be exported.
Optional Arguments
export.onnx_file
: The path to save the exported model to. The default path is in the same directory as the\*.tlt
model.export.gpu_id
: The index of the GPU that will be used to run the export. You can specify this value when the machine has multiple GPUs installed. Note that export can only run on a single GPU.
Here’s an example of using the ReIdentificationNet export command:
tao model re_identification export -e $DEFAULT_SPEC -r $RESULTS_DIR -k $KEY export.checkpoint=$TRAINED_TLT_MODEL
You can deploy the trained deep -earning and computer-vision models on edge devices–such as a Jetson Xavier,
Jetson Nano, or Tesla–or in the cloud with NVIDIA GPUs. The exported
\*.onnx
model can also be used with TAO Toolkit Triton Apps.
Running ReIdentificationNet Inference on the Triton Sample
The TAO Toolkit Triton Apps provide an inference sample for ReIdentificationNet. It consumes a TensorRT engine and supports running with a directory of query (probe) images and a directory of test (gallery) images containing the same identities.
To use this sample, you need to generate the TensorRT engine from an \*.onnx
model using
trtexec
.
Generating TensorRT Engine Using
trtexec
For instructions on generating a TensorRT engine using the trtexec
command, refer to the
trtexec guide for ReIdentificationNet.
Running the Triton Inference Sample
You can generate the TensorRT engine when starting the Triton server using the following command:
bash scripts/start_server.sh
When the server is running, you can get results from a directory of query images and a directory of test images using the following command with a client:
python tao_client.py <path_to_query_directory> \
--test_dir <path_to_test_directory>
-m re_identification_tao model \
-x 1 \
-b 16 \
--mode Re_identification \
-i https \
-u localhost:8000 \
--async \
--output_path <path_to_output_directory>
The server will perform inference on the input image directories. The results are saved as a JSON file. The following is a sample of the JSON output:
[
...,
{
"img_path": "/localhome/Data/market1501/query/1121_c3s2_156744_00.jpg",
"embedding": [-1.1530249118804932, -1.8521332740783691,..., 0.380886435508728]
},...
{
"img_path": "/localhome/Data/market1501/bounding_box_test/1377_c2s3_038007_05.jpg",
"embedding": [0.09496910870075226, 0.26107653975486755,..., 0.2835155725479126]
},...
]
End-to-End Inference Using Triton
The TAO Toolkit Triton Apps provides a sample for end-to-end inference from a directory of query images and a directory of test images. The sample downloads the Market-1501 dataset and randomly samples a subset of 100 identities. The client implicitly converts the image samples into arrays and sends them to the Triton server. The feature embedding for each image is returned and saved to the JSON output. An image of sampled matches and a figure of the CMC curve is also generated for visualization.
You can start the Triton server using the following command (only the ReIdentificationNet model will be downloaded and converted into a TensorRT engine):
bash scripts/re_id_e2e_inference/start_server.sh
Once the Triton server has started, open another terminal and use the following command to run re-identification on the query and test images using the Triton server instance that you have previously spun up:
bash scripts/re_id_e2e_inference/start_client.sh