PrunaAI
/

PrunaSuperPoint

@@ -1,172 +1,247 @@
 ---
 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](https://github.com/rpautrat/SuperPoint).
-For distillation and evaluation, we use the [UZH FPV Dataset](https://fpv.ifi.uzh.ch/) [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:
-![superpoint_structure](./static/model_structure.png)
-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 weights
-- `original-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 the `backbone_0_1` layer 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 the `16_16_24_32_64` pruning 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` - previous `pruned` model 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](http://rpg.ifi.uzh.ch/datasets/uzh-fpv-newer-versions/v3/indoor_forward_3_snapdragon_with_gt.zip). This was the trace for training the `pruned` model 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](http://rpg.ifi.uzh.ch/datasets/uzh-fpv-newer-versions/v3/outdoor_forward_1_snapdragon_with_gt.zip). 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` |
 |---|---|---|
 | ![](./static/pruned_light_indoor.png) | ![](./static/pruned_indoor.png) | ![](./static/pruned_noref_indoor.png) |
-| `pruned-light`|`pruned` |`pruned-noref` |
-|---|---|---|
 | ![](./static/pruned_light_outdoor.png) | ![](./static/pruned_outdoor.png) | ![](./static/pruned_noref_outdoor.png) |
-Alternatively, we can also compare the results with no overlay, e.g., here is the indoor image for the `pruned` model:
-![superpoint_structure](./static/pruned_indoor_separate.png)
-## 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. Includes `setup.py` script 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
 ```

 ---
 license: apache-2.0
+library_name: PrunaSuperPoint
 tags:
+- SuperPoint
+- vision
+- image-matching
+- pruned
+- pytorch
+- onnx
+pipeline_tag: keypoint-detection
 ---
+<!-- header start -->
+<!-- 200823 -->
+<div style="width: auto; margin-left: auto; margin-right: auto">
+    <a href="https://docs.pruna.ai/en/stable/index.html" target="_blank" rel="noopener noreferrer">
+        <img src="https://imgur.com/rVAgqMY.png" alt="PrunaAI" style="width: 100%; min-width: 400px; display: block; margin: auto;">
+    </a>
+</div>
+<!-- header end -->
+[![GitHub](https://img.shields.io/badge/GitHub-PrunaAI-9334E9?style=plastic&logo=github&logoColor=white)](https://github.com/PrunaAI/pruna) &nbsp;
+[![Twitter/X](https://img.shields.io/badge/Twitter%2FX-@PrunaAI-9334E9?style=plastic&logo=x&logoColor=white)](https://twitter.com/PrunaAI) &nbsp;
+[![LinkedIn](https://img.shields.io/badge/LinkedIn-PrunaAI-9334E9?style=plastic&logo=linkedin&logoColor=white)](https://www.linkedin.com/company/pruna-ai) &nbsp;
+[![Discord](https://img.shields.io/badge/Discord-Join%20us-9334E9?style=plastic&logo=discord&logoColor=white)](https://discord.com/invite/JFQmtFKCjd) &nbsp;
+[![P-Models](https://img.shields.io/badge/Performance%20Models-Try%20them%20now-9334E9?style=plastic&logo=data:image/svg%2bxml;base64,PHN2ZyB4bWxucz0iaHR0cDovL3d3dy53My5vcmcvMjAwMC9zdmciIHdpZHRoPSI0MTEiIGhlaWdodD0iNDc3IiBmaWxsPSJub25lIj48cGF0aCBmaWxsPSIjZmZmIiBkPSJNMjMuNzAzIDMwMC41MzJjLTEuODQyLS42NzQtNC4xNDItMS4wNTctNS43MDMtMi4wMjktMTIuNTU4LTcuODQ2LTguODM3LTI1LjE3MiA1LjI1NS0yOC4zNjFhLjY2LjY2IDAgMCAwIC4zNDgtLjIwMS42LjYgMCAwIDAgLjE1Ni0uMzZxNi40OS04MS45ODYgNjguNTc3LTEzNC42MjlhLjM0LjM0IDAgMCAwIC4wOTItLjEyNS4zMy4zMyAwIDAgMC0uMDI2LS4zLjMzLjMzIDAgMCAwLS4xMTItLjEwOGMtNC4xMTUtMi40OTctOC42NC00Ljc1LTEyLjM1My03LjQ0M2ExOTAgMTkwIDAgMCAxLTE4Ljc5NC0xNS41MjJjLTIuOTM2LTIuNzY4LTMuNzAzLTUuODkxLjU4LTcuNzdhLjUuNSAwIDAgMCAuMjgtLjYwOEM1MyA3NS45NSA0OC40MTcgNDguNTk5IDQ1LjUxOCAxOS42NXEtLjM5My0zLjg1Mi0uMjcxLTYuMzJhLjc1Ljc1IDAgMCAxIC4xMzQtLjQuNy43IDAgMCAxIC4zMjQtLjI1NXEyLjEwNC0uODQyIDUuODQ0LTEuNDMxYzU4LjExNC05LjEzNiAxMjcuMDA5IDguNDI1IDE0NC4zMDcgNzIuNjQ1cS4xMDMuMzY0LjE3OCAwIDUuMDk2LTI0LjA5NyA0LjYzOC00OC4xOTMtLjA0Ny0yLjc4NyAxLjg3OS00Ljc4OCAxMC40MDctMTAuOCAyNC43OTgtOC42ODdjNy40NzEgMS4wOTUgNy43NDIgNC42MzggNy4yMDkgMTEuMWEzMzQgMzM0IDAgMCAxLTExLjg4NCA2NC43MDYuMzg4LjM4OCAwIDAgMCAuMTgyLjQ0NC40LjQgMCAwIDAgLjE2NC4wNTFjOTIuNDIxIDguNDUzIDE1OC44NzUgNzkuNDg5IDE2NS4wODQgMTcxLjI3NGExIDEgMCAwIDAgLjEyOS40MjQuODMuODMgMCAwIDAgLjI5MS4yOTZjMy43NSAyLjE3IDYuMiAyLjIxNiA4Ljg1NSA2LjU5MiA1LjA2OCA4LjMzMiAxLjQ1IDIwLjQzMS04LjU5MyAyMi40N2ExLjcxIDEuNzEgMCAwIDAtMS4zNjUgMS40NjhxLTkuMTczIDc1LjM4My02NC43NjIgMTIzLjM3MWMtMjguNjIyIDI0LjcwNC02NS43NDQgMzkuNzQ5LTEwMy40NjQgNDIuMzMtMTAwLjM1OSA2Ljg2My0xODQuNDg2LTY1LjYyMi0xOTUuMDA1LTE2NS41N2EuOC44IDAgMCAwLS4xNTQtLjM5NS43NC43NCAwIDAgMC0uMzMzLS4yNU0xNTEuMzM4IDk3LjY0M2wxOS4wNjYgMTkuMzc1YS41LjUgMCAwIDAgLjE1LjEuNS41IDAgMCAwIC4xNzYuMDMxLjQzLjQzIDAgMCAwIC4zMTktLjE1cTEwLjg0Ny0xMi40NjUgNy40NDMtMjguMTU1Yy00LjUzNS0yMC44Ni0xOC41OTgtMzcuNjY0LTM3LjI3MS00OC4zODktMjEuMDY3LTEyLjY4OS00Ny44NjUtMTUuOTA1LTcyLjI4LTE0Ljg0cS0uNjE3LjAzLS4xNC40MjJsNjMuNDcyIDUzLjMyNmEuNC40IDAgMCAwIC4xNDQuMDcyLjM0LjM0IDAgMCAwIC4zMDEtLjA2OC4zLjMgMCAwIDAgLjA4OC0uMTI2IDQ3IDQ3IDAgMCAwIDMuMjgyLTEzLjA4MXEuMDIzLS4yMzMuMTI0LS40MzN0LjI2OC0uMzM0YzIuOTU1LTIuNDc4IDQuMTE1LjQ1OCA0LjY5NCAyLjg1MnEyLjMyIDkuNTM3LjE4NyAxOS4yNDRhMS40OCAxLjQ4IDAgMCAwIC40NCAxLjQwM3ptLTYuNTI2IDkuMjk1LTkuMDI0LTEwLjI3NmEuNy43IDAgMCAwLS4zNDMtLjIwNy44Ni44NiAwIDAgMC0uNDMzIDBxLTExLjY4IDMuMDIxLTIzLjExNC0uODIyYy01LjgxNi0zLjg1Mi0uNzU4LTUuNTE3IDIuNTktNS44MTZxNS45MTgtLjUzMyAxMS43MzUtMi41MDZhLjMuMyAwIDAgMCAuMTI3LS4wOC4zLjMgMCAwIDAgLjA3My0uMTMzLjM0LjM0IDAgMCAwIC4wMDEtLjE1NC4zNC4zNCAwIDAgMC0uMDctLjEzOGMtMTkuNDU5LTIxLjE1LTQwLjc2LTQwLjI0NS02Mi4yNzUtNTkuMjY0cS0uMzU2LS4zMTgtLjMxOC4xNWEzMDggMzA4IDAgMCAwIDguMzg3IDUxLjgwMmM0LjU5MiAxOC4zMDggMTMuNzU1IDM0LjE1NyAzMS44NDggNDEuODcycTI2LjI0NyAxMS4xODIgNTMuNzI5IDIuODMzYS41ODMuNTgzIDAgMCAwIC4yOTktLjg4OCAyMzUgMjM1IDAgMCAwLTEzLjIxMi0xNi4zNzNtNDUuMTI1IDQ5Ljc0NXEtMTIuNTk1IDEuMzE4LTIzLjQ0MS01LjU4M2MtMy4wNzctNC44OTkgMi4wNjYtOC41OTMgNS40NDItMTEuMTgzYS41ODQuNTg0IDAgMCAwIC4xODctLjcwMSA4LjQ2IDguNDYgMCAwIDAtMy40ODgtMy44NjIuODEuODEgMCAwIDAtLjgyMy4wMDljLTE2LjY5MSA5Ljk3Ny0zNC40ODUgMTAuNzgyLTUzLjE4NiA2Ljk2N3EtMS41Ny0uMzItMi44NTIuNjQ1LTYxLjQ1MSA0Ni4zNy02OS42NyAxMjMuOTEzLS42ODMgNi40MzMtMS44MDYgOC44MjdjLTIuOTQ1IDYuMzAyLTcuNzcgNy44NDUtMTQuMjIyIDkuNDcycS0uNDQ5LjExMyAwIC4yMzRjOS4yNzYgMi4zMSAxNC42OCA2LjE0MyAxNi4wNzQgMTYuNDAxIDguNDYyIDYyLjQ5IDQ4Ljk0IDExNy44NjQgMTEwLjM4MyAxMzguMzQxYTEwNCAxMDQgMCAwIDAgMTEuMDA2IDMuNTA3YzkyLjY0NSAyMy45MDkgMTgxLjIzMi0zMy42NjIgMjAyLjgwNC0xMjUuMDU0cTEuNjA5LTYuNzkgMy40NS0xOS4wODVjMS4xMDQtNy4zOTYgNS42OTUtMTIuMTE4IDEzLjE0Ny0xMy41NjguNDEyLS4wODQgMS45MDgtLjc1Ny4yODEtLjk5MS04LjA1MS0xLjE0MS0xMi4yNC02LjM4Ni0xMi45ODgtMTQuMTI4cS0xLjA3NS0xMS4yNS0yLjUxNi0xOS4wMzhjLTExLjc3Mi02My44OTItNTYuMzgzLTExNC41NzItMTE5LjI4NS0xMzEuNDZxLTE1LjEyLTMuODgtMzAuOTg3LTUuMjQ1YS4zODYuMzg2IDAgMCAwLS40MDIuMjUyIDE2OC42IDE2OC42IDAgMCAxLTExLjM1MiAyNC44NTRjLTQuMTMzIDcuNDA2LTguMTYzIDE1LjQ5NC0xNS43NTYgMTYuNDc2Ii8+PGNpcmNsZSBjeD0iMTQwLjA2NSIgY3k9IjI3My4wMyIgcj0iNTAuNjE4IiBzdHJva2U9IiNmZmYiIHN0cm9rZS13aWR0aD0iOS4yMDMiLz48cGF0aCBmaWxsPSIjZmZmIiBkPSJNMTQ2Ljk2NyAyNDIuMDkyYzE1LjY3MiAwIDI4LjM3NiAxMy45NjMgMjguMzc2IDMxLjE4OHMtMTIuNzA0IDMxLjE4OS0yOC4zNzYgMzEuMTg5LTI4LjM3Ny0xMy45NjQtMjguMzc3LTMxLjE4OWMwLTIuNjI1LjI5Ny01LjE3NC44NTItNy42MSAxLjYwNyAzLjcyNyA1LjMxMyA2LjMzOCA5LjYyOSA2LjMzOCA1Ljc4OSAwIDEwLjQ4MS00LjY5MyAxMC40ODItMTAuNDgycy00LjY5My0xMC40ODEtMTAuNDgyLTEwLjQ4MWMtLjc2MSAwLTEuNTA0LjA4My0yLjIxOS4yMzcgNS4xMzktNS42NzYgMTIuMjUzLTkuMTkgMjAuMTE1LTkuMTkiLz48Y2lyY2xlIGN4PSIyNjkuOTMyIiBjeT0iMjczLjAzIiByPSI1MC42MTgiIHN0cm9rZT0iI2ZmZiIgc3Ryb2tlLXdpZHRoPSI5LjIwMyIvPjxwYXRoIGZpbGw9IiNmZmYiIGQ9Ik0yNzYuODM0IDI0Mi4wOTJjMTUuNjcyIDAgMjguMzc2IDEzLjk2MyAyOC4zNzYgMzEuMTg4cy0xMi43MDQgMzEuMTg5LTI4LjM3NiAzMS4xODktMjguMzc3LTEzLjk2NC0yOC4zNzctMzEuMTg5YzAtMi42MjUuMjk3LTUuMTc0Ljg1My03LjYxIDEuNjA2IDMuNzI3IDUuMzEyIDYuMzM4IDkuNjI4IDYuMzM4IDUuNzg5IDAgMTAuNDgyLTQuNjkzIDEwLjQ4Mi0xMC40ODJzLTQuNjkzLTEwLjQ4MS0xMC40ODItMTAuNDgxYy0uNzYxIDAtMS41MDQuMDgzLTIuMjE5LjIzNyA1LjEzOS01LjY3NiAxMi4yNTQtOS4xOSAyMC4xMTUtOS4xOSIvPjxwYXRoIHN0cm9rZT0iI2ZmZiIgc3Ryb2tlLXdpZHRoPSI5LjIwMyIgZD0iTTE3Ny4xMzUgMzQ1LjMyOGMuODk5LS4yNCAyLjgwMS0uMTk3IDYuMzA0LjQ2NCA2LjMxMSAxLjE5IDE2LjU4NSA0LjIyMSAyNi4xMjIgNC4yMjEgOS40NTIgMCAxNy4xMTgtMS45MjYgMjEuMzgzLTIuNjAzLjc3LS4xMjMgMS4zODUtLjE5OSAxLjg3OC0uMjM3LS4zNCAxMy44NDMtMTIuODkzIDI3LjcxLTI3Ljg2OCAyNy43MS0xNS4xNzIgMC0yNy44NzYtMTMuODM1LTI3Ljg3Ni0yOC44MTEgMC0uMzU3LjAyOS0uNTk2LjA1Ny0uNzQ0WiIvPjwvc3ZnPg==)](https://dashboard.pruna.ai/login?utm_source=huggingface&utm_medium=org_card&utm_campaign=hf_traffic)
+<div align="center">
+<h1 style="color: #9334E9;">⚡️ PrunaSuperPoint</h1>
+<h2>TensorRT SuperPoint on Jetson Orin Nano</h2>
+<h3>
+  <span style="color: #9334E9;">~1.9× faster at 640×480</span>
+  &nbsp;·&nbsp;
+  <span style="color: #9334E9;">~2.1× faster at 1080p</span>
+  &nbsp;·&nbsp;
+  Optimized from SuperPoint
+</h3>
+</div>
+**PrunaSuperPoint is an optimization of the SuperPoint (SP) keypoint detection model. Based on [rpautrat/SuperPoint](https://github.com/rpautrat/SuperPoint) by Rémi Pautrat and Paul-Edouard Sarlin. Optimized for the TensorRT runtime on Jetson Orin Nano devices, but also applicable to other runtime backends. The ideas used can be applied to other similar network architectures.**
+This work was done in collaboration with CEA ([Aidge](https://eclipse.dev/aidge/)) for the [DeepGreen](https://deepgreen.ai/) project. You can also find a pruning tutorial for SuperPoint using Aidge operators in the [Aidge codebase](https://gitlab.eclipse.org/eclipse/aidge/aidge/-/tree/main/examples/tutorials/SuperPoint_pruning_tutorial).
+## Example
+![pruned indoor keypoints vs original, shown separately](./static/pruned_indoor_separate.png)
+## Optimization methods
+To achieve better performance, we use the following methods:
+- **Structural pruning**: It's the main optimization method used. The SuperPoint model architecture consists of eight sequential convolutional layers, and after every two convolutional layers, the model has a pooling layer, which halves the width and height of the feature map. So, the model's biggest bottlenecks are the first convolutional layers, since they have to process the image at full input resolution. By reducing the input/output channel sizes of desired convolutional layers, the runtime can be improved. In the case of channel pruning, we keep the channels corresponding to kernels with the highest mean absolute weight values, which provides a good initialization point for potential distillation training.
+- **Hierarchical top-k keypoint selection**: Instead of selecting the global top-k keypoints directly, this can be done hierarchically by splitting the image into chunks, selecting the top *k* keypoints from each chunk, and then aggregating the candidates to obtain the final result. This guarantees the original result, using the same *k* and assuming no scoring ties, while providing a noticeable speedup on constrained devices.
+- **Skipping keypoint refinement**: The base model performs two refinement steps after non-maximum suppression to recover additional high-scoring keypoints. Skipping them, even if not a direct optimization, avoids four extra full-resolution passes, which is very expensive.
+## Benchmark
+Metrics compare five different model variations.
 - `original` - base SP with original weights
+- `original-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 the `backbone_0_1` layer 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 the `16_16_24_32_64` pruning config that affects the first six layers of the model and uses our distillation checkpoint weights, which were trained on 250 examples. We use a combination of KL divergence and cross-entropy loss to recover keypoint locations, and a cosine similarity-based loss for descriptors. Previous top-k optimization is also enabled.
+- `pruned-noref` - previous `pruned` model with no NMS refinement steps.
+### Latency
+| Setting | Description |
+|---|---|
+| Model format | ONNX models serialized as TensorRT engines using the `trtexec` CLI utility. |
+| Precision | FP16 |
+| Device | Jetson Orin Nano 8GB. |
+| Metrics | 1) Mean inference time over 1000 images from the indoor trace using our benchmarking script 2) Mean inference time reported by the TensorRT `trtexec` benchmarking utility. |
+| Units | Milliseconds (ms). |
+Standard comparison:
+| Input | original | pruned | Speedup |
+|---|---:|---:|---:|
+| 640×480 | 17.8 / 16.8 | 9.6 / 8.5 | **1.85× / 1.98×** |
+| 1920×1080 | 111.1 / 109.6 | 54.2 / 52.9 | **2.05× / 2.07×** |
+All variants:
+| Input | 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 |
+### Quality
+| Setting | Description |
+|---|---|
+| Model format | PyTorch models |
+| Input resolution | Native dataset resolution of 640×480 |
+| Keypoints | 1024 keypoints per image |
+| Datasets | 1) Indoor dataset with 767 images from [UZH FPV Indoor forward facing trace #3](http://rpg.ifi.uzh.ch/datasets/uzh-fpv-newer-versions/v3/indoor_forward_3_snapdragon_with_gt.zip) 2) Outdoor dataset with 500 images from [UZH FPV Outdoor forward facing trace #1](http://rpg.ifi.uzh.ch/datasets/uzh-fpv-newer-versions/v3/outdoor_forward_1_snapdragon_with_gt.zip) |
+| Metrics | 1) Average keypoints covered, the keypoints coverage compares the spatial distribution of predictions against the original model over 8×8 image regions, with 1024 as the maximum score 2) Mean descriptor difference L2 norm, the L2 norm is calculated by comparing the descriptor vectors at shared keypoint locations and calculating the norm of their difference. |
 | 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 / - |
+<details>
+<summary>Keypoint visualization examples</summary>
+A few examples of the predicted keypoints are shown below. The first row uses `image_1_1200.png` from the Indoor trace, while the second uses `image_0_1300.png` from the Outdoor trace. Each target model is compared against the original model predictions. Green points indicate keypoints predicted at the exact same location by both models, while red and blue indicate keypoints predicted by only one of the two models.
+| `pruned-light` | `pruned` | `pruned-noref` |
 |---|---|---|
 | ![](./static/pruned_light_indoor.png) | ![](./static/pruned_indoor.png) | ![](./static/pruned_noref_indoor.png) |
 | ![](./static/pruned_light_outdoor.png) | ![](./static/pruned_outdoor.png) | ![](./static/pruned_noref_outdoor.png) |
+Although there are relatively few exact overlaps, many predictions lie very close to each other. In addition, even when target model predictions do not match the original exactly, they often still correspond to sensible keypoint locations.
+</details>
+## Quickstart
+Install the required dependencies and the package in editable mode:
+```bash
+pip install -r requirements.txt
+pip install -e .
+```
+Download the default dataset, **UZH FPV Indoor forward-facing trace #3**, to `src/superpoint_pruning/distillation/data`:
+```bash
+superpoint-pruning setup
+```
+Compare the keypoint predictions from the pruned model against the original SuperPoint model on an image:
+```bash
+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
 ```
+The resulting comparison is written to `./test.png`. To run the distillation and evaluation scripts, see the [Annex: Run your own distillation and compare inference time](#annex-run-your-own-distillation-and-compare-inference-time) section.
+## Limitations
+- Other popular SuperPoint implementations exist, especially versions paired with matchers such as LightGlue. This repository does not use them because of their restrictive licenses. The work is meant to show that SuperPoint-like models can be pruned and distilled. You can apply the same methods to those implementations after the necessary structural changes and a distillation run with your configs.
+- The right way to judge a pruned SuperPoint 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 currently looking at a specific task.
+- Distillation used 250 frames from [UZH-FPV Indoor forward facing trace #3](http://rpg.ifi.uzh.ch/datasets/uzh-fpv-newer-versions/v3/indoor_forward_3_snapdragon_with_gt.zip) (767 eval images). Direct training images were excluded from evaluation, but results may still be biased because the frames are sequential video frames. Outdoor forward-facing trace #1 was not used for recovery training.
+- Skipping NMS refinement usually still produces sensible keypoints, but that depends on the task.
+## License
+This repository is licensed under the **Apache License 2.0**.
+### Third-party Attributions
+- **SuperPoint implementation and weights** — Rémi Pautrat and Paul-Edouard Sarlin, [rpautrat/SuperPoint](https://github.com/rpautrat/SuperPoint).
+- **SuperPoint method** — Daniel DeTone, Tomasz Malisiewicz, and Andrew Rabinovich, *SuperPoint: Self-Supervised Interest Point Detection and Description*, CVPR Workshops, 2018.
+- **UZH-FPV dataset** — Jeffrey Delmerico, Titus Cieslewski, Henri Rebecq, Matthias Faessler, and Davide Scaramuzza, *Are We Ready for Autonomous Drone Racing? The UZH-FPV Drone Racing Dataset*, ICRA, 2019. [Dataset](https://fpv.ifi.uzh.ch/)
+## What's next?
+- Experiment with PrunaSuperPoint for computer vision, drone, and edge-device workloads
+- Compress your own models with [Pruna](https://github.com/PrunaAI/pruna) and give us a ⭐️
+## Annex: Run your own distillation and compare inference time
+For installing the package with dependencies for running the basic distillation, evaluation, and export scripts, run:
+```bash
 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:
+```bash
 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:
+```bash
 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:
+```bash
 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 comparison plot between a target model and the original on a specific image, one can for example use:
+```bash
 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:
+```bash
 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:
+```bash
 superpoint-pruning benchmark --model-path superpoint.engine
 ```