license: apache-2.0
library_name: superpoint-pruned
tags:
- SuperPoint
1. SuperPoint Pruned
This repository focuses on optimizing the SuperPoint (SP) keypoint detection model for the TensorRT runtime on Jetson Orin Nano devices. To achieve better performance, we use methods like structural pruning, distillation training for performance recovery, and general model improvements for serialized model exporting. Despite the focus on the Jetson devices, the optimization methods are general enough to likely benefit other runtime backends as well and the ideas used can be applied to other similar network architectures.
As a base we use the open source PyTorch implementation of Superpoint and corresponding base weights by Rémi Pautrat and Paul-Edouard Sarlin, which can be found in this repository. For distillation and evaluation, we use the UZH FPV Dataset [1].
Note that there are also other popular SuperPoint implementations available, especially versions paired with keypoint matching models such as LightGlue. In this repository we won't be using them because of their restrictive license and instead our work is intended more as a proof of concept to demonstrate feasibility of optimizing SuperPoint-like models. One can in theory apply the methods covered here to alternative implementations by making necessary structural changes and running the distillation pipeline with desired configurations.
[1] J. Delmerico, T. Cieslewski, H. Rebecq, M. Faessler, and D. Scaramuzza, “Are We Ready for Autonomous Drone Racing? The UZH-FPV Drone Racing Dataset,” in IEEE International Conference on Robotics and Automation (ICRA), 2019.
2. Optimization methods
The following provides a brief overview of the optimization methods used in the repository.
Pruning
Pruning is the main optimization method used and main motivation for this repository. For context, the network backbone of the default SuperPoint model architecture consists of eight sequential convolutional layers as shown in the figure:
The main observation is that after every two convolutional layers, the model has a pooling layer, which halves the width and height of the feature map. This means that the biggest bottlenecks of the model are the first convolutional layers, since they have to process the image at full input resolution. Hence the runtime can be improved by reducing the input/output channel sizes of desired convolutional layers, which our pruning interface allows to do with a simple class method addition. Note that in case of channel pruning, as a default we keep the channels that correspond to kernels with the highest mean absolute weight values, which is a good heuristic to retain as much of original performance as possible and also offers a good initialization point for potential distillation training.
In the benchmarking section we outline two pruning configurations.
One prunes only the backbone_0_1 layer input channels from 64 to 32 (which means also adjusting backbone_0_0 output channels accordingly) with no additional training, since this is the most expensive layer in the model.
Another config we explore adjusts the first six layers of the model and because of that it requires distillation training to recover accuracy.
We provide the training checkpoint under src/superpoint_pruning/weights/16_16_24_32_64.ckpt and we use a combination of KL divergence and cross entropy loss for recovering the keypoint locations and a cosine similarity based loss for the descriptors.
Here the filename 16_16_24_32_64 refers to the new values of input channels for the corresponding layers starting from backbone_0_1, backbone_1_0 and so on.
Hierarchical top-k keypoint selection
The original model extracts the final k keypoints by choosing the top scoring ones out of all possible pixel locations of the input image. Since the k value is normally much smaller than the number of candidate, then this can be done hierarchically by first splitting the image into chunks and finding the top k keypoints in each chunk, after which the final results can be aggregated. Using the same k and assuming no scoring ties, we are guaranteed the original result while getting a small speedup which is noticeable for constrained devices.
Skipping keypoint refinement
This is not a direct optimization, but rather a useful observation in case the model runtime is especially important. The base model uses non-maximal supression to find top keypoints in each image region. While doing this, the original implementation also does two additional refinement steps, which allows to extract other high scoring keypoints that were supressed by the first pass. This usually results in more fine-grained keypoints in the final output, but in total the model does four extra passes over the full image resolution, which is very expensive. Hence we include an option to skip this refinement, as the outputs still stay relatively sensible without it, but this depends on the use case.
3. Benchmark results
In the following we compare the quality and inference time of five different model variations:
original- base SP with original weightsoriginal-topk- base SP with hierarchical top-k enabled. Each variant uses the original k for sub-chunks, which guarantees the original results. The resolution 640x480 uses 32 chunks and 1920x1080 uses 36 chunks.pruned-light- base SP where input channels of thebackbone_0_1layer have been pruned from 64 to 32 with no additional training. This demonstrates the effect of the most expensive layer and how the kernel selection heuristic can retain performance. Previous top-k optimization is also enabled.pruned- uses the16_16_24_32_64pruning config that affects the first six layers of the model (as described in the methods section) and uses our distillation checkpoint weights, which is trained on 250 examples and recovers the performance. Exact training configuration is referenced in the usage section. Previous top-k optimization is also enabled.pruned-noref- previousprunedmodel with no NMS refinement steps to demonstrate the effects of it.
Quality evaluation
The previous models were evaluated using PyTorch using the native dataset image size of 640x480 as input and 1024 keypoints in each case. We differentiate two specific dataset traces:
- Indoor - 767 images from UZH FPV Indoor forward facing trace #3. This was the trace for training the
prunedmodel and while we removed the direct images that were used for training, these results are still biased because of the sequential nature of the video frames. - Outdoor - 500 images from the UZH FPV Outdoor forward facing trace #1. Images from this trace were not used for performance recovery training of the pruned models and should serve as an example for the generalization capability of the distillation.
We report two scores for each combination in the table below as average number of keypoints covered / mean descriptor difference L2 norm. The average number of keypoints covered score splits the original image into 8x8 chunks and checks if the target model predicted the same amount of keypoints in each chunk as the original. This isn't a perfect score but it also isn't as restrictive as exact match comparison and gives a general proxy score to assess, if the target model's keypoints have similar image coverage to the original (with 1024 being the best results, covering all original regions). In this case, a lower score doesn't mean that keypoints are definitely bad, but they might be more concentrated to a specific region than the original. The L2 norm is calculated by comparing the descriptor vectors at shared keypoint locations and calculating the norm of their difference. This score serves just as a quantitative signal for relative performance - in an ideal case this could be combined with a matching model to measure if the target model descriptor vectors are still recognizable.
Note that generally, the best way to assess the pruned versions of SP is to test it in the intended downstream task and see if it still retains performance. We don't explore it here since we are not looking at a specific task currently.
| Trace | original | original-topk | pruned-light | pruned | pruned-noref |
|---|---|---|---|---|---|
| Inside | 1024 / 0 | 1024 / 0 | 637 / 0,92 | 817 / 0,24 | 700 / - |
| Outside | 1024 / 0 | 1024 / 0 | 485 / 0,89 | 754 / 0,27 | 671 / - |
In the following you can also see a few visualized examples of the predicted keypoints. The used models from left to right are pruned-light, pruned and pruned-noref respectively. The first set uses image_1_1200.png from the Indoor trace and the second set uses image_0_1300.png from the outdoor trace.
All of the target model keypoints are compared to the original model predictions, hence the green points are keypoints which both models predicted exactly, while red and blue are points which one predicted, but the other did not. Notice that while there aren't many exactly overlapping points, there are many points where the original and target model predictions lie very close. Additionally, while not all of the target model prediction match the original, then often the predicted points are still sensible keypoint locations.
pruned-light |
pruned |
pruned-noref |
|---|---|---|
 |
 |
 |
pruned-light |
pruned |
pruned-noref |
|---|---|---|
 |
 |
 |
Alternatively, we can also compare the results with no overlay, e.g., here is the indoor image for the pruned model:
Inference time evaluation
For this section, all of the previous model configurations were exported as ONNX models and then serialized as TensorRT models using the trtexec CLI utility (with TRT version 10.3.0) with the --fp16 flag to enable fp16 precision, which offers additional speedup for no quality degradation.
Testing is done on a Jetson Orin Nano 8GB version and two times are reported: the first is the mean inference time across 1000 images from the previous indoor trace using our python benchmarking script and the second is the mean inference time reported by the TensorRT benchmarking utility (used as trtexec --loadEngine=serialized_engine_file --useCudaGraph --useSpinWait).
All times are reported as milliseconds.
| Input resolution | original | original-topk | pruned-light | pruned | pruned-noref |
|---|---|---|---|---|---|
| 640 × 480 | 17,8 / 16,8 | 17,0 / 16,0 | 14,4 / 13,5 | 9,6 / 8,5 | 8,3 / 7,0 |
| 1920 × 1080 | 111,1 / 109,6 | 102,5 / 102,1 | 86,9 / 84,6 | 54,2 / 52,9 | 45,1 / 43,5 |
4. Project structure
The code is organized into separate modules under src/superpoint_pruning/:
export.py- utility script for exporting a pruned SuperPoint model as ONNX.models/superpoint.py- base optimized SuperPoint PyTorch model.distillation/- PyTorch Lightning trainer script and utilities for distillation training. Includessetup.pyscript for downloading the base dataset that was used for training.evaluation/- quality evaluation scripts and TensorRT wrapper for inference timing measurements.weights/- base SuperPoint weights and precomputer distillation checkpoints.
5. Setup and usage
For installing the package with dependencies for running the basic distillation, evaluation, and export scripts, run:
pip install -r requirements.txt
pip install -e .
To install the base dataset images (UZH FPV Indoor forward facing trace #3) used for training to the default path src/superpoint_pruning/distillation/data, run:
superpoint-pruning setup
Alternatively, specify the desired trace download link with the --download-url flag, e.g., the one for the UZH FPV Outdoor forward facing trace #1, and the local dataset directory name with --output-dir-name.
Run distillation training
The base PyTorch Lightning training config that we used to produce the pruned model can be found under src/superpoint_pruning/distillation/base_config.yaml.
Since this includes the hyperparameters and the data setup used for our training, then this should allow to reproduce the trained checkpoint that is provided under src/superpoint_pruning/weights/16_16_24_32_64.ckpt. The training script itself is very basic but can serve as a starting point for further experiments.
To replicate the training, assuming that the main Indoor data trace is installed in the previous, run:
superpoint-pruning setup --generate-gt
superpoint-pruning train
The first command computes and presaves the target keypoint and descriptor feature maps and the other initiate the Lightning training script.
Run evaluation
The first row of quality evaluation results can be reproduced using the following commands corresponding to the different target models:
superpoint-pruning evaluate
superpoint-pruning evaluate --hierarchical
superpoint-pruning evaluate --backbone_0_1 32 --hierarchical
superpoint-pruning evaluate --pruning-config ./src/superpoint_pruning/weights/16_16_24_32_64.ckpt --hierarchical
superpoint-pruning evaluate --pruning-config ./src/superpoint_pruning/weights/16_16_24_32_64.ckpt --hierarchical --skip-refinement
Note that in this script and other similar ones, we can define the pruning either with format --backbone_x_y new_input_channel_size or by pointing it to a ckpt file named in the format a_b_...ckpt, which is translated to --backbone_0_1 a, --backbone_1_0 b and so on.
For the evaluation using the outdoor trace (or some other trace), the --image-dir parameter has to be specified including the other specific options to use the same image range (--no-skip --start-idx 1000 --end-idx 1500).
Plot keypoint comparisons
For a quick keypoint prediction comparsions plot between a target model and the original on a specific image, one can for example use:
superpoint-pruning plot --image-name image_1_1200.png --output ./test.png --num-keypoints 512 --pruning-config ./src/superpoint_pruning/weights/16_16_24_32_64.ckpt --overlay
This once again defaults to the UZH FPV Indoor forward facing trace #3 and accepts file names from the dataset.
Export model to ONNX
To export the desired PyTorch model to ONNX:
superpoint-pruning export --backbone_0_1 32
The model pruning configuration can once again be controlled with same flags as shown in the evaluation example.
As an example, we have published the ONNX export of the original model to https://huggingface.co/EclipseAidge/SuperPoint.
Run TRT inference benchmark
In case the device supports TensorRT (with the necessary packages installed) and there is a serialized TRT model (here a dummy example of superpoint.engine), then the local python benchmark with a TRT wrapper can be run as:
superpoint-pruning benchmark --model-path superpoint.engine