trending/spirula-studio
1
0
Fork 0
mirror of https://github.com/harry7557558/spirula-studio.git synced 2026-10-02 02:44:54 +08:00
Code Issues Packages Projects Releases Wiki Activity

aliked + lightglue for sfm

Browse Source
This commit is contained in:
Harry Chen
2026-08-03 14:16:45 -04:00
parent 70e06ed0df
commit 8bac858fe9
39 changed files with 5333 additions and 59 deletions
Download Patch File Download Diff File Expand all files Collapse all files
cmake/SsplatNn.cmake
+28 -2
View File
@@ -141,6 +141,29 @@ target_compile_options(ssplat_sam PRIVATE
$<$<COMPILE_LANGUAGE:CXX>:${SPLAT_CXX_FLAGS}>)
set_property(TARGET ssplat_sam PROPERTY CXX_STANDARD 17)
# ---------------------------------------------------------------------------
# ssplat_aliked -- the ALIKED / LightGlue learned SfM frontend
#
# Its checkpoints are ONNX files fetched from COLMAP's releases and parsed in
# process (docs/notes/aliked-port-plan.md), so this needs no protobuf, no
# onnxruntime and no converter -- src/aliked/model/Onnx.cpp is a varint walk.
#
# The ops it needs that nn/ did not have -- deformable convolution, point-wise
# grid sample, average pooling, row L2 normalization -- are general and went
# into nn/shaders. What is here is only what could not be general: the
# detector's suppression rule and soft-argmax, and its coordinate conversions.
# ---------------------------------------------------------------------------
ssplat_nn_shaders(aliked ${SSPLAT_SRC}/aliked/shaders SSPLAT_ALIKED_EMBED)
file(GLOB_RECURSE SSPLAT_ALIKED_SOURCES CONFIGURE_DEPENDS ${SSPLAT_SRC}/aliked/*.cpp)
list(FILTER SSPLAT_ALIKED_SOURCES EXCLUDE REGEX "/tests/")
add_library(ssplat_aliked STATIC ${SSPLAT_ALIKED_SOURCES} ${SSPLAT_ALIKED_EMBED})
target_link_libraries(ssplat_aliked PUBLIC ssplat_nn)
target_compile_options(ssplat_aliked PRIVATE
$<$<COMPILE_LANGUAGE:CXX>:${SPLAT_CXX_FLAGS}>)
set_property(TARGET ssplat_aliked PROPERTY CXX_STANDARD 17)
# ---------------------------------------------------------------------------
# ssplat_video -- container demux + VK_KHR_video_decode_*
#
@@ -165,11 +188,14 @@ endif()
# Tests -- one executable per file, as in src/backend/tests/
# ---------------------------------------------------------------------------
file(GLOB SSPLAT_NN_TESTS CONFIGURE_DEPENDS
${SSPLAT_SRC}/nn/tests/*.cpp ${SSPLAT_SRC}/sam/tests/*.cpp)
${SSPLAT_SRC}/nn/tests/*.cpp ${SSPLAT_SRC}/sam/tests/*.cpp
${SSPLAT_SRC}/aliked/tests/*.cpp)
foreach(test_src ${SSPLAT_NN_TESTS})
get_filename_component(test_name ${test_src} NAME_WE)
add_executable(${test_name} ${test_src})
target_link_libraries(${test_name} PRIVATE ssplat_sam)
# Every test links every library above it: the three are small, and one
# rule here beats a per-directory list that drifts.
target_link_libraries(${test_name} PRIVATE ssplat_sam ssplat_aliked)
set_property(TARGET ${test_name} PROPERTY CXX_STANDARD 17)
target_compile_options(${test_name} PRIVATE
$<$<COMPILE_LANGUAGE:CXX>:${SPLAT_CXX_FLAGS}>)
cmake/SsplatSfm.cmake
+12
View File
@@ -151,6 +151,18 @@ add_library(ssplat_sfm STATIC
)
target_include_directories(ssplat_sfm PUBLIC ${SSPLAT_SRC})
target_link_libraries(ssplat_sfm PUBLIC Vulkan::Vulkan Threads::Threads)
# The learned frontend (src/aliked/) is optional: it sits on the inference
# layer, which is SSPLAT_BUILD_SAM. Without it `--features aliked-*` is a
# usage error that says so, and nothing else changes -- SfM keeps building on
# a machine that only wants SIFT. PUBLIC because sfm/feature/Extractor.h's
# factory is compiled into whatever links this.
if(SSPLAT_BUILD_SAM)
target_link_libraries(ssplat_sfm PUBLIC ssplat_aliked)
target_compile_definitions(ssplat_sfm PUBLIC SSPLAT_HAVE_ALIKED=1)
else()
target_compile_definitions(ssplat_sfm PUBLIC SSPLAT_HAVE_ALIKED=0)
endif()
target_compile_options(ssplat_sfm PRIVATE
$<$<COMPILE_LANGUAGE:CXX>:${SPLAT_CXX_FLAGS}>)
set_property(TARGET ssplat_sfm PROPERTY CXX_STANDARD 17)
docs/notes/aliked-port-plan.md
+269
View File
@@ -0,0 +1,269 @@
# ALIKED + LightGlue: the learned frontend
The plan for `src/sfm/README.md`'s unstarted item 15, and the record of what
was decided and why. Companion to [sfm-port-plan.md](sfm-port-plan.md) and
[segmentation-port.md](segmentation-port.md).
Goal: `--features aliked-n16rot|aliked-n32` and `--matcher lightglue` as
alternatives to GPU SIFT and the brute-force matcher, reaching the same
`FeatureSet` / `IFeatureMatcher` contracts, with **no new dependency** — no
onnxruntime, no PyTorch, no self-hosted weights.
## What is actually being ported
Nothing from COLMAP's `src/colmap/feature/aliked.cc`. That file is 341 lines of
ONNX plumbing around a black box, and `bruteforce-matcher.onnx` is a **5 KB**
graph doing cosine + ratio + cross-check, which `sfm/feature/Matcher.h` already
does. What gets ported is:
1. **Two networks, reimplemented on `src/nn/`** — the layer that exists for
exactly this ("A learned feature detector for SfM … goes on top of it
unchanged"). This is the work.
2. **COLMAP's pre/post-processing conventions, literally.** Image
normalization, the pad-to-/32 replicate padder, and the keypoint coordinate
un-normalization. These are five lines each and every one of them is a place
to lose half a pixel silently.
3. **COLMAP's option semantics**, which are not the same shape as ours:
`min_cossim = 0.85`, `max_ratio = 1.0` — i.e. the ratio test is *off* and an
absolute similarity threshold is the only filter. Our `reduce()` has no such
threshold today.
## Weights: the ONNX file is the checkpoint
**We do not host or convert anything.** `AlikedWeights` fetches the same
artifact COLMAP does, from the same URL, verifies the same SHA-256, and parses
it in process:
```
https://github.com/colmap/colmap/releases/download/3.13.0/aliked-n16rot.onnx
39c423d0a6f03d39ec89d3d1d61853765c2fb6a8b8381376c703e5758778a547 3.0 MB
https://github.com/colmap/colmap/releases/download/3.13.0/aliked-n32.onnx
a077728a02d2de1a775c66df6de8cfeb7c6b51ca57572c64c680131c988c8b3c 4.2 MB
https://github.com/colmap/colmap/releases/download/3.13.0/aliked-lightglue.onnx
b9a5de7204648b18a8cf5dcac819f9d30de1a5961ef03756803c8b86c2dceb8d
```
(The URI triples live in COLMAP's `src/colmap/feature/resources.h`; ours mirror
them.) Consequences, all good:
- Bit-identical weights to what COLMAP runs, so a parity gate against
`colmap feature_extractor --FeatureExtraction.type ALIKED_N16ROT` means
something.
- No converter to keep in sync, nothing to publish, nothing to commit. The
cache directory and the download path are `src/app/gui/ModelCache.cpp`'s
business, unchanged.
- Licensing is the easy case: ALIKED is BSD-3-Clause and LightGlue is
Apache-2.0, so this is a `needs_tick = false` family, unlike SAM 3.
### Parsing ONNX without a dependency
An ONNX file is a protobuf. We need **initializers only** — the graph structure
is hard-coded in our forward pass — so the reader is a ~200-line varint walk of
three nested messages and no schema:
```
ModelProto field 7 -> GraphProto
GraphProto field 5 -> repeated TensorProto (the initializers)
TensorProto field 1 dims (packed or repeated), 2 data_type,
8 name, 9 raw_data, 4 float_data
```
Everything else in the file is skipped by wire type. Verified against both
checkpoints: 2519 nodes, **69 initializers, all f32**, 2.71 MB of a 3.00 MB
file. `src/aliked/model/Onnx.cpp` is that walk and nothing more.
### What the checkpoints say
The two variants have **identical graphs** — 2519 nodes, 69 initializers each —
and differ only in `M`, the number of SDDH sample positions, which is
`desc_head.agg_weights.shape[0]`. So: one code path, `M` read from the file,
`n16rot` vs `n32` is a URL. `n16rot` is architecturally `n16`, trained
rotation-robust.
```
block1 ConvBlock(3->16) conv1[16,3,3,3] bn1 conv2[16,16,3,3] bn2
block2 ResBlock(16->32) + downsample[32,16,1,1]+bias
block3 ResBlock(32->64) conv{1,2} are DeformableConv2d:
offset_conv[18,Cin,3,3]+bias, regular_conv[Co,Cin,3,3]
block4 ResBlock(64->128) same shape
conv1..4 [32,{16,32,64,128},1,1] the dim/4 projections, no bias
score_head.{0,2,4,6} [8,128,1,1] [4,8,3,3] [4,4,3,3] [1,4,3,3], no bias
desc_head.offset_conv.0 [2M,128,3,3]+bias (K=3 patch, VALID)
desc_head.offset_conv.2 [2M,2M,1,1]+bias
desc_head.sf_conv [128,128,1,1], no bias
desc_head.agg_weights [M,128,128]
hw_grid [25,2] the 5x5 soft-argmax grid
```
`offset_conv` out-channels are `2*k*k = 18`, not `3*k*k` — **no modulation
mask**, so the deformable convolution is the plain torchvision
`deform_conv2d(x, offset, w)` with `mask = None`. BatchNorm is *not* folded in
the export (8 `BatchNormalization` nodes, `running_*` present as initializers);
we fold it into the preceding conv at load, which is why the loader wants the
conv weight and its BN together.
## Module layout
```
src/aliked/ # a sibling of src/sam/, on top of nn/
├── README.md
├── Aliked.h # the public surface: AlikedExtractor, LightGlueMatcher
├── model/Onnx.{h,cpp} # the protobuf walk
├── model/Weights.{h,cpp} # host tensors -> VRAM, BN folding, hparams from shapes
├── model/AlikedModel.cpp # backbone, score head, DKD, SDDH
├── model/LightGlue.{h,cpp} # 9 self/cross layers + assignment
├── shaders/aliked.slang # SDDH, NMS, soft-argmax, the assignment tail
└── tests/aliked_test.cpp
src/sfm/feature/Extractor.{h,cpp} # IFeatureExtractor + factory (SIFT is one impl)
src/sfm/feature/LearnedMatcher.{h,cpp} # LightGlue behind IFeatureMatcher + factory
```
(As built. Neither sfm header includes an `aliked/` one, so `src/sfm/` still
compiles without the inference layer; the factories say so at run time.)
`ssplat_aliked` goes in `cmake/SsplatNn.cmake` next to `ssplat_sam`, with its
own `ssplat_nn_shaders(aliked …)` edge and its own declare/ensure pair — *a
static library whose only content is a static initializer is not linked*, so a
new shader directory that skips this comes back "no shader module". `SsplatSfm`
links it when `SSPLAT_BUILD_SAM AND SSPLAT_BUILD_SFM`.
Model-specific constants stay here, never in `nn/`.
## The nn/ op gap
Confirmed against the exported graph's op mix (`Conv, Selu, Sigmoid,
AveragePool, Resize, MaxPool, GridSample, TopK, Einsum, BatchNormalization`).
| op | status |
|---|---|
| `Act::Selu` | new — ALIKED's gate is `nn.SELU` everywhere. One enum row, one `apply_act` case. |
| `avgpool` (2×2 s2, 4×4 s4) | new; `maxpool2x2` exists, this is the same shape |
| `resize_bilinear` **align_corners=True** | ALIKED uses `nn.Upsample(align_corners=True)`; ours is hard-wired to torch's `align_corners=False` mapping. Add a flag; **do not** change the default, mask upsampling depends on it. |
| **deformable conv 3×3** | new, and the only substantial kernel. offset conv (existing `conv2d`, Cout=2·9) → clamp to ±max(H,W)/4 → `deform_im2col` feeding the existing chunked im2col+GEMM. Same chunking as `OpConv.cpp`, bilinear fetch instead of a direct one, **zero outside the image** (torchvision's rule, not clamp-to-edge). |
| **`grid_sample_points`** | new. Bilinear sample of a `[H,W,C]` map at N arbitrary positions → `[N,C]`. SDDH needs it; `roi_align` is close but not a substitute. |
| NMS + compaction | new, but it is `sfm/feature/Sift.h`'s pattern verbatim: 5×5 max-pool equality (`simple_nms(radius=2)`, twice), append through an atomic counter, read the count back, exact top-K on the host. A device sort is not worth it at a few thousand candidates. |
| soft-argmax | new, tiny — 5×5 patch, T=0.1, the `hw_grid` initializer is its coordinate table |
| **SDDH** | new. One workgroup per keypoint: gather the 3×3×C patch → two small matmuls → clamp → M bilinear samples → `sf_conv` → `einsum('ncp,pcd->nd')` over `agg_weights[M,C,C]` → L2 normalize. ~1 GFLOP at N=4k, M=16, C=128. |
| **LightGlue: nothing** | `linear` (bias/residual/act), `layer_norm`, `attention` (4 heads × 64), `rope` — our freqs layout `[n, hd/2, 2]` *is* LightGlue's rotary form. Only the tail needs a kernel: log-double-softmax over N₀×N₁ + mutual-NN + threshold. |
Every one lands with an `nn_ops_test` case against an independent scalar CPU
reference. That file is why the model layer can treat this API as trustworthy.
## sfm/ plumbing
1. **`IFeatureExtractor`.** `extractDirectory` news `SiftExtractor` directly
today; a factory on the feature type also keeps sfm's Vulkan device
uncreated when ALIKED is selected — which matters, see (7).
2. **`FeatureSet`**: `dtype = F32, dim = 128`. Already allowed (D1). ~1 MB per
image at 2048 keypoints, the same order as SIFT at 8192.
3. **Persist the detection score.** ALIKED has no scale and no orientation.
`PairSelection::topScaleSubset` ranks by `scale`, so with an all-zero scale
it silently picks an arbitrary 512 keypoints — a wrong answer that looks
like a working feature. `features.bin` goes to **v5** with a per-keypoint
score (`Keypoint::response`, which exists and is unwritten), and the subset
ranks by scale-or-score.
4. **Working resolution.** COLMAP's `EffMaxImageSize()` is **1600 for ALIKED**
against 3200 for SIFT, and it means it: the aggregated map is 128 channels
at *full* resolution, so 1600×1200 is ~1 GB fp32 for one tensor. `--quality`
gets a per-frontend default.
5. **RGB in, `/255`, no mean/std.** `GrayImage::rgb` already exists behind
`want_color = true`, which `extract` already sets for point colors.
6. **The coordinate chain**, copied from `aliked.cc` exactly: pad to /32 by
replicating the right and bottom edges; keypoints come back normalized to
[-1,1] over the *padded* extent; `px = (nx + 1) * 0.5 * (padded_w - 1) +
0.5`; drop anything outside the original bounds. ALIKED's origin is the
top-left pixel *center*, ours is its corner — that is what the `+0.5` is.
7. **Two Vulkan devices.** `nn::Device` and `sfm::VkContext` are separate
contexts. Under `match` with LightGlue, pair selection (sfm) and LightGlue
(nn) would both be live; sequence them rather than interleaving. Converging
is sfm-port-plan phase 6 and does not block this.
8. **Matching.** Brute force stays the existing kernel: quantize the
L2-normalized descriptors to int8 *on upload* (128 B/descriptor, exactly
what the packed-dot path and the residency budget are sized for), keeping
F32 on disk for LightGlue. Add the absolute-similarity threshold that
COLMAP's defaults require. LightGlue is ~100+ GFLOP for a 2048×2048 pair —
tens of milliseconds — so it runs **only behind pair selection**, never over
a raw exhaustive list.
9. **Config**: `--features`, `--matcher`, `--aliked-max-features`,
`--aliked-min-score`, `--lightglue-min-score`, `--aliked-model`. One
`SFM_CONFIG_FIELDS` row each; the GUI's editor gets them for free.
## Phases, each with a gate
| | work | gate | |
|---|---|---|---|
| P0 | ONNX reader + weight store | every expected tensor present with the expected shape; `aliked_test` over both checkpoints | **done** |
| P1 | the new nn ops | `nn_ops_test` cases vs scalar CPU references | **done** |
| P2 | backbone + score head | score-map cosine vs an ORT reference dump on one image. 80% of the bugs die here. | **done** |
| P3 | DKD + SDDH | vs COLMAP's own extractor on ~20 images: keypoint repeatability, descriptor cosine ≥ 0.99 | **done** |
| P4 | sfm plumbing, brute force end to end | `tools/sfm/eval_poses.py` AUC against the SIFT baseline | **done** |
| P5 | LightGlue | same gate; the win should show on wide-baseline and low-texture captures | **done** |
`aliked-n32` was free after P3, as predicted: one constant read from the file,
no code change, and it hit parity on the first run. `aliked-t16` and the ONNX
brute-force matcher are deliberately not ported.
### What P0-P3 measured
`tools/aliked/compare_colmap.py` against `colmap feature_extractor
--FeatureExtraction.type ALIKED_N16ROT`, four images, both variants:
98.6-99.8% of keypoints within 1 px, mean offset below 0.003 px in both axes,
descriptor cosine 0.99995 mean. The residual is at the score threshold — peaks
within float noise of `min_score` survive on one side and not the other.
56-66 ms at ~1.1 MP against ~700 ms for COLMAP's onnxruntime CPU path. The
table lives in `src/aliked/README.md`.
### What P4-P5 measured
`ssplat sfm auto` on a 20-image wide-baseline capture, exhaustive matching, one
camera group. Registration and point count, not just match volume:
| frontend + matcher | features/img | pairs kept | inliers | registered | points |
|---|---|---|---|---|---|
| SIFT + brute force | 8192 | 68/190 | 7 703 | 55% | — |
| ALIKED + brute force | 4096 | 190/190 | 24 597 | 100% | 6 795 |
| ALIKED + LightGlue | 4096 | 75/75 (shortlist) | 52 981 | 100% | 17 343 |
LightGlue is ~53 ms for a 2029 x 2044 pair against ~3.5 ms for brute force, so
its preset switches `--pairs` to `prefilter`; exhaustive matching with it is
hours where pair selection is minutes. It matched ORT exactly on a real pair --
702 matches both sides, 100% identical partners, max score difference 0.024.
**COLMAP's matching defaults were measured and rejected.** Its ALIKED
brute-force settings are `min_cossim 0.85` with the ratio test off; on this
data an absolute 0.85 cosine rejects ~70% of mutual-nearest matches (their
median cosine is 0.726) and left 24/190 pairs. What the data wanted instead was
the ratio test kept and *loosened* to 0.92 -- a learned descriptor's second-best
distance sits much closer to its best (median ratio 0.826, just the wrong side
of SIFT's 0.8), which is the same property that makes LightGlue worth having.
`--min-similarity` still exists for anyone who wants COLMAP's shape of test.
Two things in this plan were wrong, both found by bisecting against the graph
rather than by reading the reference implementation, and both worth recording
because they are the shape of bug this port produces:
- **`ResBlock`'s second gate is outside the residual add**, not inside it.
The plan asserted the opposite on the strength of a Selu *count* (17), which
is consistent with both placements. Gating inside gave a detector that still
found 16% of the same keypoints and descriptors that were uncorrelated --
i.e. it looked like a working feature extractor. The exported graph's
dataflow (`/block2/gate_1/Selu` consumes `/block2/Add_output_0`) settles it,
and a node's *inputs* are worth more than its count.
- **A GEMM cannot be run in place.** `linear(feat, feat, w)` for the descriptor
head's `sf_conv` races: every output column re-reads the whole input row.
Elementwise ops in `nn/` do alias safely and `LinearOpts::residual` is
explicitly allowed to, which is what made it look permissible.
## Risks, named up front
- **The deformable conv's sampling rule.** torchvision zero-fills outside the
image and orders offsets `(dy, dx)` per tap; getting either wrong produces a
score map that looks plausible and matches badly.
- **`grid_sample` conventions.** `align_corners` and the ±1 normalization
appear three times (upsample, deform conv, SDDH) and are not the same
convention in all three.
- **The coordinate chain** in (6): four frames — normalized, padded pixels,
extraction pixels, source pixels — and `scaleKeypoints` is a fifth hop.
- **fp16.** SAM's policy is f16 for matmul/conv weights. ALIKED is 3 MB total;
keep everything f32 until P3 passes, then measure. The activations are what
cost memory here, not the weights.
src/aliked/Aliked.h
+87
View File
@@ -0,0 +1,87 @@
#pragma once
// ALIKED keypoint extraction: an RGB image in, keypoints and 128-D float
// descriptors out.
//
// The public surface of src/aliked/. Everything below it -- the ONNX reader,
// the weight store, the forward pass -- is an implementation detail, and
// nothing here exposes Vulkan, nn::Tensor or the checkpoint's layout, so
// src/sfm/ can hold one of these without depending on the inference layer's
// headers.
//
// Conventions, all inherited from COLMAP's extractor so that a reconstruction
// built on our features and one built on COLMAP's are comparable
// (docs/notes/aliked-port-plan.md):
//
// * Input is 8-bit interleaved RGB, normalized by 1/255 and nothing else.
// * The image is padded to a multiple of 32 by replicating its right and
// bottom edges; the network needs it, and keypoints found in the padding
// are dropped.
// * Output coordinates are in the ORIGINAL image's pixels with the top-left
// pixel's *corner* at (0, 0) -- COLMAP's convention, half a pixel off
// ALIKED's own, which puts the top-left pixel's centre there.
// * Descriptors are L2-normalized floats, so a dot product is a cosine.
#include <cstdint>
#include <string>
#include <vector>
namespace aliked {
struct ExtractOptions {
// Top-scoring keypoints kept. COLMAP's AlikedExtractionOptions default.
int max_num_features = 2048;
// Minimum detection score, applied to the sub-pixel-refined score exactly
// where COLMAP applies it.
float min_score = 0.2f;
// DKD's NMS radius and soft-argmax temperature. Not exposed as flags
// anywhere; they are the trained detector's constants, here so the
// forward pass has one place to read them from.
int nms_radius = 2;
float temperature = 0.1f;
// Candidate list capacity. Every local maximum of the score map is
// collected before the top-K cut, and a 1600x1200 image has tens of
// thousands; saturating warns rather than failing, as GPU SIFT's lists do.
uint32_t max_candidates = 262144;
};
struct Keypoint {
float x = 0, y = 0; // original-image pixels, corner origin
float score = 0;
};
struct Features {
int width = 0, height = 0; // the image the coordinates refer to
int desc_dim = 0;
std::vector<Keypoint> keypoints;
std::vector<float> descriptors; // keypoints.size() * desc_dim
};
// Owns a checkpoint and the working buffers for one image at a time.
// Not thread-safe: one Extractor per thread, and they share the process-wide
// inference device.
class Extractor {
public:
Extractor();
~Extractor();
Extractor(const Extractor&) = delete;
Extractor& operator=(const Extractor&) = delete;
// `model` is a known id ("aliked-n16rot", "aliked-n32"), which is fetched
// and cached, or a path to an .onnx file. Throws nn::Error on anything
// that is not an ALIKED checkpoint.
void load(const std::string& model);
bool loaded() const;
// Descriptor width of the loaded checkpoint (128 for every released one).
int descriptorDim() const;
// `rgb` is width*height*3 interleaved bytes.
Features extract(const uint8_t* rgb, int width, int height,
const ExtractOptions& opts = {});
private:
struct Impl;
Impl* impl_ = nullptr;
};
} // namespace aliked
src/aliked/Common.h
+25
View File
@@ -0,0 +1,25 @@
#pragma once
// The names aliked/ borrows from the inference layer (src/nn/), exactly as
// sam/Common.h does for src/sam/. Include this at the top of an aliked/
// translation unit instead of qualifying nn:: at every call site.
//
// Nothing model-specific belongs here -- that is what aliked/Aliked.h is for.
#include "nn/core/Error.h"
#include "nn/core/Half.h"
#include "nn/core/Log.h"
#include "nn/core/Parallel.h"
namespace nn { namespace vk {} }
namespace aliked {
namespace vk = ::nn::vk;
using ::nn::Error;
using ::nn::fail;
using ::nn::log_level;
using ::nn::now_ms;
using ::nn::ScopedTimer;
using ::nn::parallel_for;
} // namespace aliked
src/aliked/README.md
+168
View File
@@ -0,0 +1,168 @@
# ALIKED (`src/aliked/`)
The learned SfM frontend: ALIKED keypoints and 128-D float descriptors, on the
inference layer (`src/nn/`), with no onnxruntime, no PyTorch and no converter.
Status: **done and wired in.** The extractor and LightGlue both match COLMAP's
ONNX implementations, and `ssplat sfm --features aliked-n16rot --matcher
lightglue` runs the whole pipeline.
## Rules
- **Vulkan only**, like `src/sfm/` and `src/sam/`. Built with
`SSPLAT_BUILD_SAM` (which is what builds `ssplat_nn`).
- **Model-specific things live here, never in `nn/`.** The general ops this
needed — deformable convolution, point-wise grid sample, average pooling, row
L2 normalization, SELU — went into `nn/` and are tested there. What is left
in `shaders/aliked.slang` is only what could not be general: the detector's
suppression rule, its soft-argmax, and two coordinate conversions.
- **The checkpoint is COLMAP's, fetched from COLMAP's release URL and parsed in
process.** We host nothing and convert nothing.
## The checkpoint is an ONNX file
`model/Onnx.cpp` is a ~300-line varint walk of protobuf that reads
*initializers only* — the graph structure is hard-coded in `AlikedModel.cpp`,
so nothing has to understand a node, an operator or a shape rule.
The point of this is parity. Both sides run the same bytes, so a difference
against `colmap feature_extractor --FeatureExtraction.type ALIKED_N16ROT` is
ours and can be bisected, which is exactly how the two bugs below were found.
It also means there is no export to keep in step with an upstream re-export.
`aliked-n16rot` and `aliked-n32` have **byte-identical graph structure** — 2519
nodes, 69 initializers each — and differ only in `M`, the number of SDDH sample
positions, which is `desc_head.agg_weights.shape[0]`. So there is one code
path, `M` comes from the file, and the variant is a URL. `n16rot` is
architecturally `n16`, trained rotation-robust.
BatchNorm is *not* folded in the export (8 `BatchNormalization` nodes,
`running_*` present as initializers); `model/Weights.cpp` folds each into the
conv before it, with the epsilon read from the node rather than assumed.
## Layout
```
Aliked.h the public surface: Extractor, ExtractOptions, Features
Common.h the names borrowed from nn/, as sam/Common.h does
model/Onnx.{h,cpp} the protobuf walk
model/Fetch.{h,cpp} URL + SHA-256 table, cache, download through curl
model/Weights.{h,cpp} initializers -> device tensors; BN folding; hparams
model/AlikedModel.cpp the forward pass, read against nets/aliked.py
model/LightGlue.{h,cpp} the matcher: 9 self/cross layers + one assignment
shaders/aliked.slang NMS, soft-argmax, the SDDH coordinate conversions,
and LightGlue's assignment reduction
tests/aliked_test.cpp checkpoint shapes, SHA-256, extraction and matching
```
On the SfM side, `sfm/feature/Extractor.h` and `sfm/feature/LearnedMatcher.h`
are the seams this plugs into; neither includes an `aliked/` header, so
`src/sfm/` still builds without the inference layer and says so at run time.
## Five conventions that are not guessable
Each of these was wrong at some point, and each failure looked like a working
network rather than a broken one.
1. **`ResBlock`'s second gate is outside the residual add** —
`gate(bn2(conv2(gate(bn1(conv1(x))))) + downsample(x))`. Reading the module
source suggests the gate sits inside the branch; the exported graph settles
it (`/block2/gate_1/Selu` takes `/block2/Add_output_0`). Gating inside gave
16% keypoint agreement instead of 99.6%, with uncorrelated descriptors.
2. **Two offset layouts in one network.** The deformable convolutions use
torchvision's `(dy, dx)` interleaved per tap. The descriptor head's offset
conv emits `2M` channels that ALIKED reads as `view(N, 2, M)` — the first
`M` are all the x components, the second `M` all the y. Swapping either
produces plausible-looking descriptors that match badly.
3. **`align_corners=True`** on every upsample and every normalized coordinate
downstream of one — which is *not* the convention `nn::resize_bilinear`
defaults to, because mask upsampling needs the other one.
4. **LightGlue's fused qkv is `[head][dim][3]`** — q, k and v interleaved per
element, which no stride can express. The projection's output rows are
permuted at load so it becomes `[3][head][dim]`, the layout
`nn::attention`'s q/k/v strides already address.
5. **LightGlue's assignment score contains `sim` twice**, once from the row
log-softmax and once from the column one. The arg-max is over
`2*sim[i][j] - lse_row[i] - lse_col[j] + ls0[i] + ls1[j]`; dropping the
factor changes which column wins, not just the value, and pushes every
score far enough below zero that nothing matches at all.
## Memory
The aggregated feature map is `dim` (128) channels at **full resolution**, so
it dominates everything: ~1 GB at COLMAP's default working size for this
extractor (1600 px, against 3200 for SIFT). `plan_arena_bytes` sizes the whole
forward pass up front because the arena will not grow while anything is live.
Keeping that map in f16 would halve it and is the obvious next move; it is not
done because `l2_normalize_rows` would need an f16 output path, i.e. a second
numeric path under a kernel that currently has one.
## Testing
```bash
./build/aliked_test # cached checkpoints, or SKIP
./build/aliked_test --fetch # download from COLMAP's releases
./build/aliked_test --image IMG.jpg --out /tmp/ours.bin
./build/aliked_test --match /tmp/a.bin /tmp/b.bin # LightGlue on two dumps
```
The checkpoint gate is strict about *shapes* — we do not own these weights and
cannot embed a golden copy, so what it can check is that every tensor the
forward pass will ask for exists at the width the rest of the model assumes.
The gate that matters is parity against COLMAP, which needs a COLMAP built with
ONNX support:
```bash
colmap feature_extractor --database_path /tmp/db.db --image_path IMAGES \
--FeatureExtraction.type ALIKED_N16ROT
./build/aliked_test --image IMAGES/x.jpg --out /tmp/ours.bin
python3 tools/aliked/compare_colmap.py /tmp/db.db /tmp/ours.bin
```
Measured on four images, `n16rot` and `n32`, on an RTX 5070 Laptop:
| | keypoints (ours / COLMAP) | matched ≤1 px | mean offset | descriptor cosine |
|---|---|---|---|---|
| n16rot | 2029 / 2029 | 99.6% | (+0.0002, −0.0001) px | 0.999983 |
| n16rot | 2044 / 2044 | 99.8% | (−0.0004, −0.0002) px | 0.999982 |
| n16rot | 1071 / 1079 | 99.2% | (+0.0022, −0.0010) px | 0.999953 |
| n16rot | 1541 / 1529 | 98.6% | (+0.0023, +0.0005) px | 0.999971 |
| n32 | 2048 / 2048 | 99.8% | (−0.0002, −0.0002) px | 0.999978 |
The residual disagreement is at the score threshold: peaks within float noise
of `min_score` survive on one side and not the other. Extraction is 56–66 ms at
~1.1 MP, against ~700 ms for COLMAP's onnxruntime CPU path on the same image.
`nn_ops_test`'s "Learned frontend" section covers the ops this needed, each
against an independent scalar CPU reference.
## LightGlue
Nine transformer layers, each self-attention within an image then
cross-attention between the two, and one assignment head. It needed **no new
general ops**: `linear`, `layer_norm`, `attention` and `rope` cover the forward
pass, and our RoPE frequency layout already *is* LightGlue's rotary encoding.
The assignment tail is the only kernel, and it exists to avoid materializing a
second N0 x N1 matrix (16 MB at 2048 keypoints each). Every term of the score
separates into a row constant and a column constant, so the two log-sum-exps
are vectors and the arg-max applies them on the fly -- three passes over `sim`
instead of five.
The export runs the assignment on the last layer only: LightGlue's early exit
and token pruning are not in the graph, so they are not implemented. Both are
speedups, not behaviour.
Measured against ORT on a real pair: **702 matches on both sides, 100%
identical partners**, max score difference 0.024 (fp32 accumulation order over
nine layers). 53 ms for a 2029 x 2044 pair.
## Not done yet
1. **f16 for the aggregated map**, per "Memory" above.
2. LightGlue's early exit and width pruning, worth roughly 2x and absent from
the exported graph.
3. The candidate list is collected and top-K'd on the host. That is the same
shape as GPU SIFT's stage boundary and has not been measured as a cost.
src/aliked/model/AlikedModel.cpp
+560
View File
@@ -0,0 +1,560 @@
// ALIKED's forward pass, next to its PyTorch reference.
//
// Reading order matches the network: encoder, feature aggregation, score head,
// DKD (detect), SDDH (describe). There is no computation graph -- each op runs
// when called and writes into an arena tensor -- so this file is meant to be
// read against nets/aliked.py and nets/blocks.py line by line.
//
// The three conventions worth knowing before reading anything else, because
// each is a place where a plausible guess is wrong:
//
// * The encoder runs at FULL resolution and aggregates 128 channels there.
// That is what makes the working resolution (1600 px, COLMAP's default for
// ALIKED against 3200 for SIFT) a memory decision and not a speed one.
// * The deformable convolutions' offsets are (dy, dx) interleaved per tap --
// torchvision's layout. The descriptor head's offsets are the same
// network's, and are grouped [x0..xM-1, y0..yM-1] instead. Two conventions
// in one model; see sddh_positions in shaders/aliked.slang.
// * Upsampling is align_corners=True, and so is every normalized coordinate
// downstream of it.
#include "aliked/Aliked.h"
#include "aliked/Common.h"
#include "aliked/model/Fetch.h"
#include "aliked/model/Weights.h"
#include "nn/Ops.h"
#include "nn/Tensor.h"
#include "nn/vk/Memory.h"
#include "nn/vk/Stream.h"
#include "nn/vk/EmbeddedSpirv.h"
#include <algorithm>
#include <cmath>
#include <cstring>
#include <numeric>
#include <vector>
// This library's SPIR-V blobs. Registered by an explicit call from load(),
// not by a static initializer: ssplat_aliked is a static archive, and an
// object nothing references is not linked -- the kernels would come back "no
// shader module". See nn/vk/EmbeddedSpirv.h.
NN_DECLARE_EMBEDDED_MODULES(aliked)
namespace aliked {
namespace {
using nn::Act;
using nn::ConvOpts;
using nn::DType;
using nn::LinearOpts;
using nn::Tensor;
constexpr int kPadDivisor = 32; // the encoder's total downsampling factor
int64_t round_up(int64_t v, int64_t m) { return (v + m - 1) / m * m; }
// The arena is a bump allocator that refuses to grow while anything is live,
// so the whole forward pass has to be sized before it starts.
//
// This is also the honest statement of what ALIKED costs: the aggregated map
// is `dim` channels at FULL resolution, and at COLMAP's default working size
// for this extractor (1600 px) that one tensor is ~1 GB. Everything else is
// small next to it. Halving it by keeping the aggregate in f16 is the obvious
// next move and is not done yet -- the row normalization would need an f16
// output path, which is a second numeric path under a kernel that currently
// has one.
int64_t plan_arena_bytes(int64_t Hp, int64_t Wp, const AlikedHparams& hp) {
const int64_t P = Hp * Wp;
auto at = [&](int64_t div, int64_t ch) { return P / (div * div) * ch; };
// What stays live from the start of the encoder to the end of the
// descriptor head.
int64_t live = P * 3 // the padded input
+ P * hp.c1 // x1
+ at(2, hp.c1) // p2
+ at(2, hp.c2) // x2
+ at(8, hp.c2) // p3
+ at(8, hp.c3) // x3
+ at(32, hp.c3) // p4
+ at(32, hp.c4) // x4
+ P * hp.dim // the aggregate
+ P; // the score map
// The largest scoped transient on top of that. Aggregation is the peak: a
// full-resolution projection plus its upsampled copy. The encoder blocks
// (2 x their own output) and the score head (16 channels) are both below
// it, and the NMS masks (2 ints per pixel) are too.
const int64_t transient = 2 * P * hp.dim4;
// The convolution column workspace (nn::conv2d chunks to ~32 MiB) plus
// slack for the detector's lists and the descriptor head's small matrices,
// whose size depends on the keypoint count rather than the image.
return (live + transient) * 4 + (96ll << 20);
}
} // namespace
struct Extractor::Impl {
AlikedWeights weights;
vk::Arena arena{"aliked"};
// ---- encoder pieces -------------------------------------------------
// gate(bn2(conv2(gate(bn1(conv1(x)))))). BatchNorm was folded into each
// conv's weight and bias at load, so both convs carry a bias here and
// neither has one in the checkpoint.
//
// Every block writes into a caller-allocated `out` and scopes its own
// temporaries. That is not a style choice: the arena is a bump allocator,
// so a temporary returned past its block would stay live for the rest of
// the forward pass -- at full resolution that is 64 MB per block that
// nothing reads again.
void convBlock(const Tensor& out, const Tensor& x, const char* prefix,
int64_t cout) {
const int64_t H = x.shape[0], W = x.shape[1];
vk::ArenaScope scope(arena);
Tensor h = nn::arena_tensor(arena, DType::F32, H, W, cout);
conv3x3(h, x, weights.getf("%s.conv1.weight", prefix),
weights.getf("%s.conv1.bias", prefix), Act::Selu);
conv3x3(out, h, weights.getf("%s.conv2.weight", prefix),
weights.getf("%s.conv2.bias", prefix), Act::Selu);
}
// gate(bn2(conv2(gate(bn1(conv1(x))))) + downsample(x)).
//
// The second gate is OUTSIDE the residual add. That is torchvision's
// BasicBlock shape, and it is not what a reading of the module source
// suggests -- the exported graph settles it: /block2/gate_1/Selu takes
// /block2/Add_output_0 as its input. Gating before the add instead cost a
// 6x drop in keypoint agreement with the reference and left the
// descriptors uncorrelated, while still looking like a working detector.
void resBlock(const Tensor& out, const Tensor& x, const char* prefix, int64_t cout,
bool deformable) {
const int64_t H = x.shape[0], W = x.shape[1];
const float max_offset = 0.25f * (float)std::max(H, W);
vk::ArenaScope scope(arena);
Tensor h = nn::arena_tensor(arena, DType::F32, H, W, cout);
if (deformable)
deform3x3(h, x, prefix, 1, max_offset, Act::Selu);
else
conv3x3(h, x, weights.getf("%s.conv1.weight", prefix),
weights.getf("%s.conv1.bias", prefix), Act::Selu);
if (deformable)
deform3x3(out, h, prefix, 2, max_offset, Act::None);
else
conv3x3(out, h, weights.getf("%s.conv2.weight", prefix),
weights.getf("%s.conv2.bias", prefix), Act::None);
// The 1x1 shortcut projection, which every ResBlock here has (they all
// change width).
Tensor id = nn::arena_tensor(arena, DType::F32, H, W, cout);
ConvOpts o;
o.bias = weights.getf("%s.downsample.bias", prefix);
nn::conv2d(arena, id, x, weights.getf("%s.downsample.weight", prefix), 1, 1, o);
nn::add(out.view(H * W, cout), out.view(H * W, cout), id.view(H * W, cout), 1.0f,
1.0f, Act::Selu);
}
void conv3x3(const Tensor& out, const Tensor& in, const Tensor& w, const Tensor& bias,
Act act) {
ConvOpts o;
o.pad_y = o.pad_x = 1;
o.bias = bias;
o.act = act;
nn::conv2d(arena, out, in, w, 3, 3, o);
}
// offset = offset_conv(x).clamp(+-max(h,w)/4); deform_conv2d(x, offset, w).
void deform3x3(const Tensor& out, const Tensor& in, const char* prefix, int idx,
float max_offset, Act act) {
const int64_t H = in.shape[0], W = in.shape[1];
vk::ArenaScope scope(arena);
Tensor offset = nn::arena_tensor(arena, DType::F32, H, W, 2 * 3 * 3);
ConvOpts oo;
oo.pad_y = oo.pad_x = 1;
oo.bias = weights.getf("%s.conv%d.offset_conv.bias", prefix, idx);
nn::conv2d(arena, offset, in,
weights.getf("%s.conv%d.offset_conv.weight", prefix, idx), 3, 3, oo);
ConvOpts co;
co.pad_y = co.pad_x = 1;
co.bias = weights.getf("%s.conv%d.regular_conv.bias", prefix, idx);
co.act = act;
nn::deform_conv2d(arena, out, in, offset,
weights.getf("%s.conv%d.regular_conv.weight", prefix, idx), 3, 3,
max_offset, co);
}
};
// ---------------------------------------------------------------------------
Extractor::Extractor() : impl_(new Impl) {}
Extractor::~Extractor() { delete impl_; }
bool Extractor::loaded() const { return impl_->weights.loaded(); }
int Extractor::descriptorDim() const { return impl_->weights.hparams().desc_dim; }
void Extractor::load(const std::string& model) {
NN_ENSURE_EMBEDDED_MODULES(aliked);
impl_->weights.load(resolve_model(model));
}
Features Extractor::extract(const uint8_t* rgb, int width, int height,
const ExtractOptions& opts) {
NN_CHECK(impl_->weights.loaded(), "Extractor::extract before load()");
NN_CHECK(rgb != nullptr && width > 0 && height > 0, "Extractor::extract: empty image");
Impl& im = *impl_;
const AlikedHparams& hp = im.weights.hparams();
vk::Arena& arena = im.arena;
Features out;
out.width = width;
out.height = height;
out.desc_dim = hp.desc_dim;
const int64_t Wp = round_up(width, kPadDivisor);
const int64_t Hp = round_up(height, kPadDivisor);
NN_CHECK(Hp / 32 >= 1 && Wp / 32 >= 1, "image too small for ALIKED");
arena.reserve((uint64_t)plan_arena_bytes(Hp, Wp, hp));
vk::ArenaScope root(arena);
// ---- input: /255, padded to a multiple of 32 by edge replication ----
//
// COLMAP's InputPadder, done on the host for the same reason it does it
// there: it is one pass over the image next to a backbone that is hundreds,
// and it keeps the exact replicate-right-and-bottom rule visible.
Tensor img = nn::arena_tensor(arena, DType::F32, Hp, Wp, 3);
{
std::vector<float> host((size_t)Hp * Wp * 3);
for (int64_t y = 0; y < Hp; ++y) {
const int64_t sy = std::min<int64_t>(y, height - 1);
for (int64_t x = 0; x < Wp; ++x) {
const int64_t sx = std::min<int64_t>(x, width - 1);
const uint8_t* src = rgb + ((size_t)sy * width + sx) * 3;
float* dst = host.data() + ((size_t)y * Wp + x) * 3;
dst[0] = src[0] * (1.0f / 255.0f);
dst[1] = src[1] * (1.0f / 255.0f);
dst[2] = src[2] * (1.0f / 255.0f);
}
}
nn::tensor_from_host(img, host.data(), (int64_t)host.size());
}
// ---- encoder ----
// x1 full res, x2 at 1/2, x3 at 1/8, x4 at 1/32 -- which is where the
// pad-to-32 requirement comes from.
Tensor x1 = nn::arena_tensor(arena, DType::F32, Hp, Wp, hp.c1);
im.convBlock(x1, img, "block1", hp.c1);
Tensor p2 = nn::arena_tensor(arena, DType::F32, Hp / 2, Wp / 2, hp.c1);
nn::avgpool(p2, x1, 2);
Tensor x2 = nn::arena_tensor(arena, DType::F32, Hp / 2, Wp / 2, hp.c2);
im.resBlock(x2, p2, "block2", hp.c2, /*deformable=*/false);
Tensor p3 = nn::arena_tensor(arena, DType::F32, Hp / 8, Wp / 8, hp.c2);
nn::avgpool(p3, x2, 4);
Tensor x3 = nn::arena_tensor(arena, DType::F32, Hp / 8, Wp / 8, hp.c3);
im.resBlock(x3, p3, "block3", hp.c3, /*deformable=*/true);
Tensor p4 = nn::arena_tensor(arena, DType::F32, Hp / 32, Wp / 32, hp.c3);
nn::avgpool(p4, x3, 4);
Tensor x4 = nn::arena_tensor(arena, DType::F32, Hp / 32, Wp / 32, hp.c4);
im.resBlock(x4, p4, "block4", hp.c4, /*deformable=*/true);
// ---- aggregation: gate(conv_i(x_i)) upsampled to full res, concatenated ----
Tensor x1234 = nn::arena_tensor(arena, DType::F32, Hp, Wp, hp.dim);
{
const Tensor* blocks[4] = {&x1, &x2, &x3, &x4};
for (int i = 0; i < 4; ++i) {
vk::ArenaScope scope(arena);
const Tensor& src = *blocks[i];
const int64_t h = src.shape[0], w = src.shape[1];
Tensor proj = nn::arena_tensor(arena, DType::F32, h, w, hp.dim4);
ConvOpts o;
o.act = Act::Selu; // the 4 projections' gates -- see the Selu count
nn::conv2d(arena, proj, src, im.weights.getf("conv%d.weight", i + 1), 1, 1, o);
const Tensor* full = &proj;
Tensor up;
if (h != Hp || w != Wp) {
up = nn::arena_tensor(arena, DType::F32, Hp, Wp, hp.dim4);
nn::resize_bilinear(up, proj, /*align_corners=*/true);
full = &up;
}
// torch.cat along channels, as a strided write into the slot this
// level owns.
nn::strided_copy(x1234.offsetElems(i * hp.dim4), full->view(Hp * Wp, hp.dim4),
Hp * Wp, hp.dim4, hp.dim4, hp.dim);
}
}
// ---- score head: sigmoid(conv(...)) over the UNnormalized aggregate ----
// The explicit ndim=3 matters: a trailing 1 would otherwise be inferred
// away and the single-channel map would arrive at conv2d as [H, W].
Tensor score = nn::arena_tensor(arena, DType::F32, Hp, Wp, 1, 1, /*ndim=*/3);
{
vk::ArenaScope scope(arena);
Tensor s0 = nn::arena_tensor(arena, DType::F32, Hp, Wp, 8);
ConvOpts o;
o.act = Act::Selu;
nn::conv2d(arena, s0, x1234, im.weights.get("score_head.0.weight"), 1, 1, o);
Tensor s1 = nn::arena_tensor(arena, DType::F32, Hp, Wp, 4);
im.conv3x3(s1, s0, im.weights.get("score_head.2.weight"), {}, Act::Selu);
Tensor s2 = nn::arena_tensor(arena, DType::F32, Hp, Wp, 4);
im.conv3x3(s2, s1, im.weights.get("score_head.4.weight"), {}, Act::Selu);
im.conv3x3(score, s2, im.weights.get("score_head.6.weight"), {}, Act::Sigmoid);
}
// The descriptor head reads the CHANNEL-normalized aggregate; the score
// head above read it raw. In place, so the 128-channel full-resolution map
// exists once.
nn::l2_normalize_rows(x1234.view(Hp * Wp, hp.dim), x1234.view(Hp * Wp, hp.dim));
// ---- DKD: NMS, top-K, sub-pixel refinement ----
struct NmsParams {
uint64_t mask, supp, scores;
uint32_t H, W;
int32_t radius;
uint32_t groups_per_row;
};
struct CollectParams {
uint64_t out, counter, mask, scores;
uint32_t H, W;
int32_t border;
float min_score;
uint32_t cap, groups_per_row;
};
std::vector<uint32_t> cand;
uint32_t n_cand = 0;
{
vk::ArenaScope scope(arena);
const int64_t px = Hp * Wp;
Tensor mask = nn::arena_tensor(arena, DType::I32, px);
Tensor supp = nn::arena_tensor(arena, DType::I32, px);
Tensor list = nn::arena_tensor(arena, DType::I32, (int64_t)opts.max_candidates, 3);
Tensor counter = nn::arena_tensor(arena, DType::I32, 4);
vk::Stream::get().zero(counter.ptr, 16);
NmsParams np{};
np.mask = mask.ptr;
np.supp = supp.ptr;
np.scores = score.ptr;
np.H = (uint32_t)Hp;
np.W = (uint32_t)Wp;
np.radius = opts.nms_radius;
const vk::SpecList spec{0u, 0u};
vk::Stream::get().dispatchFlat("aliked.nms_init", spec, px, 256, &np, sizeof(np),
&np.groups_per_row);
for (int round = 0; round < 2; ++round) {
vk::Stream::get().dispatchFlat("aliked.nms_suppress", spec, px, 256, &np,
sizeof(np), &np.groups_per_row);
vk::Stream::get().dispatchFlat("aliked.nms_recover", spec, px, 256, &np,
sizeof(np), &np.groups_per_row);
}
CollectParams cp{};
cp.out = list.ptr;
cp.counter = counter.ptr;
cp.mask = mask.ptr;
cp.scores = score.ptr;
cp.H = (uint32_t)Hp;
cp.W = (uint32_t)Wp;
cp.border = opts.nms_radius;
cp.min_score = opts.min_score;
cp.cap = opts.max_candidates;
vk::Stream::get().dispatchFlat("aliked.nms_collect", spec, px, 256, &cp,
sizeof(cp), &cp.groups_per_row);
vk::Stream::get().download(&n_cand, counter.ptr, 4);
if (n_cand > opts.max_candidates) {
NN_LOG_WARN("[aliked] candidate list saturated (%u > %u); raise "
"max_candidates\n",
n_cand, opts.max_candidates);
n_cand = opts.max_candidates;
}
if (n_cand) {
cand.resize((size_t)n_cand * 3);
vk::Stream::get().download(cand.data(), list.ptr, (uint64_t)n_cand * 12);
}
}
if (n_cand == 0) return out;
// Top-K by the peak's own score, which is what DKD's top_k mode ranks by.
// The tie-break is by position, for the same reason GPU SIFT's is (D16):
// the GPU appended these through an atomic, so their order varies run to
// run, and an unstable order here would make the whole reconstruction
// irreproducible.
std::vector<uint32_t> order((size_t)n_cand);
std::iota(order.begin(), order.end(), 0u);
auto score_of = [&](uint32_t i) {
float s;
std::memcpy(&s, &cand[(size_t)i * 3 + 2], 4);
return s;
};
const uint32_t keep =
std::min<uint32_t>(n_cand, opts.max_num_features > 0
? (uint32_t)opts.max_num_features
: n_cand);
auto better = [&](uint32_t a, uint32_t b) {
const float sa = score_of(a), sb = score_of(b);
if (sa != sb) return sa > sb;
if (cand[(size_t)a * 3 + 1] != cand[(size_t)b * 3 + 1])
return cand[(size_t)a * 3 + 1] < cand[(size_t)b * 3 + 1];
return cand[(size_t)a * 3] < cand[(size_t)b * 3];
};
std::partial_sort(order.begin(), order.begin() + keep, order.end(), better);
order.resize(keep);
std::vector<int32_t> centers((size_t)keep * 2);
for (uint32_t i = 0; i < keep; ++i) {
centers[(size_t)i * 2] = (int32_t)cand[(size_t)order[i] * 3];
centers[(size_t)i * 2 + 1] = (int32_t)cand[(size_t)order[i] * 3 + 1];
}
struct RefineParams {
uint64_t out_xy, out_score, centers, scores;
uint32_t H, W, N;
int32_t radius;
float temperature;
uint32_t groups_per_row;
};
struct CenterParams {
uint64_t out, xy;
uint32_t N, groups_per_row;
};
struct SddhPosParams {
uint64_t out, xy, offsets;
uint32_t N, M, W, H;
float max_offset;
uint32_t groups_per_row;
};
const int64_t N = keep;
const int64_t M = hp.M, K = hp.K, D = hp.desc_dim;
std::vector<float> host_xy((size_t)N * 2), host_score((size_t)N);
std::vector<float> host_desc((size_t)N * D);
{
vk::ArenaScope scope(arena);
Tensor tc = nn::arena_tensor(arena, DType::I32, N, 2);
vk::Stream::get().upload(tc.ptr, centers.data(), (uint64_t)N * 8);
Tensor xy = nn::arena_tensor(arena, DType::F32, N, 2);
Tensor kscore = nn::arena_tensor(arena, DType::F32, N);
RefineParams rp{};
rp.out_xy = xy.ptr;
rp.out_score = kscore.ptr;
rp.centers = tc.ptr;
rp.scores = score.ptr;
rp.H = (uint32_t)Hp;
rp.W = (uint32_t)Wp;
rp.N = (uint32_t)N;
rp.radius = opts.nms_radius;
rp.temperature = opts.temperature;
const vk::SpecList spec{0u, 0u};
vk::Stream::get().dispatchFlat("aliked.refine_keypoints", spec, N, 256, &rp,
sizeof(rp), &rp.groups_per_row);
// ---- SDDH ----
// The patch centre is the truncation of the REFINED position, not the
// peak the refinement started from; they differ by a pixel often
// enough to matter.
Tensor ic = nn::arena_tensor(arena, DType::I32, N, 2);
CenterParams cp{};
cp.out = ic.ptr;
cp.xy = xy.ptr;
cp.N = (uint32_t)N;
vk::Stream::get().dispatchFlat("aliked.to_int_centers", spec, N * 2, 256, &cp,
sizeof(cp), &cp.groups_per_row);
Tensor patches = nn::arena_tensor(arena, DType::F32, N, hp.dim * K * K);
nn::patch_gather(patches, x1234, ic, (int)K);
Tensor off1 = nn::arena_tensor(arena, DType::F32, N, 2 * M);
LinearOpts l1;
l1.bias = im.weights.get("desc_head.offset_conv.0.bias");
l1.act = Act::Selu;
nn::linear(off1, patches, im.weights.get("desc_head.offset_conv.0.weight"), l1);
Tensor off2 = nn::arena_tensor(arena, DType::F32, N, 2 * M);
LinearOpts l2;
l2.bias = im.weights.get("desc_head.offset_conv.2.bias");
nn::linear(off2, off1, im.weights.get("desc_head.offset_conv.2.weight"), l2);
Tensor pos = nn::arena_tensor(arena, DType::F32, N * M, 2);
SddhPosParams sp{};
sp.out = pos.ptr;
sp.xy = xy.ptr;
sp.offsets = off2.ptr;
sp.N = (uint32_t)N;
sp.M = (uint32_t)M;
sp.W = (uint32_t)Wp;
sp.H = (uint32_t)Hp;
sp.max_offset = 0.25f * (float)std::max(Hp, Wp);
vk::Stream::get().dispatchFlat("aliked.sddh_positions", spec, N * M, 256, &sp,
sizeof(sp), &sp.groups_per_row);
Tensor feat = nn::arena_tensor(arena, DType::F32, N * M, hp.dim);
nn::grid_sample_points(feat, x1234, pos, /*align_corners=*/true);
// sf_conv is a 1x1 conv over the channel axis at every (keypoint,
// position), i.e. one GEMM over N*M rows. No bias in the checkpoint.
//
// A separate output, NOT in place: every output column re-reads the
// whole input row, so a GEMM writing over its own input races with
// itself. (Elementwise ops here do alias safely, and `residual` below
// is explicitly allowed to -- one thread per element. A matmul is the
// case where that reasoning does not carry.)
Tensor sfeat = nn::arena_tensor(arena, DType::F32, N * M, hp.dim);
LinearOpts lsf;
lsf.act = Act::Selu;
nn::linear(sfeat, feat, im.weights.get("desc_head.sf_conv.weight"), lsf);
// einsum('ncp,pcd->nd') as M accumulating matmuls. Each reads the
// p-th position's row out of the [N, M, C] buffer, which is what
// x_row_stride is for -- no gather, no permute.
Tensor descs = nn::arena_tensor(arena, DType::F32, N, D);
const Tensor agg = im.weights.get("desc_head.agg_weights_t");
for (int64_t p = 0; p < M; ++p) {
LinearOpts la;
la.x_row_stride = M * hp.dim;
if (p > 0) la.residual = descs;
nn::linear(descs, sfeat.offsetElems(p * hp.dim).view(N, hp.dim),
agg.slice0(p, 1).view(D, hp.dim), la);
}
nn::l2_normalize_rows(descs, descs);
vk::Stream::get().download(host_xy.data(), xy.ptr, (uint64_t)N * 8);
vk::Stream::get().download(host_score.data(), kscore.ptr, (uint64_t)N * 4);
vk::Stream::get().download(host_desc.data(), descs.ptr, (uint64_t)N * D * 4);
}
// ---- to COLMAP's frame, dropping the padding and the weak ----
//
// ALIKED puts the top-left pixel's CENTRE at (0, 0); COLMAP (and this
// repository) put its CORNER there, hence the +0.5. Keypoints found in the
// replicated padding are outside the original bounds and go.
out.keypoints.reserve((size_t)N);
out.descriptors.reserve((size_t)N * D);
for (int64_t i = 0; i < N; ++i) {
if (host_score[(size_t)i] < opts.min_score) continue;
const float px = host_xy[(size_t)i * 2] + 0.5f;
const float py = host_xy[(size_t)i * 2 + 1] + 0.5f;
if (px < 0.0f || px >= (float)width || py < 0.0f || py >= (float)height) continue;
out.keypoints.push_back({px, py, host_score[(size_t)i]});
out.descriptors.insert(out.descriptors.end(), host_desc.begin() + (size_t)i * D,
host_desc.begin() + (size_t)(i + 1) * D);
}
NN_LOG_DEBUG("[aliked] %dx%d (padded %lldx%lld): %u candidates -> %lld -> %zu\n",
width, height, (long long)Wp, (long long)Hp, n_cand, (long long)N,
out.keypoints.size());
return out;
}
} // namespace aliked
src/aliked/model/Fetch.cpp
+243
View File
@@ -0,0 +1,243 @@
#include "aliked/model/Fetch.h"
#include "nn/core/Error.h"
#include "nn/core/Log.h"
#include <cstdio>
#include <cstdlib>
#include <cstring>
#include <filesystem>
#include <fstream>
#include <vector>
namespace fs = std::filesystem;
namespace aliked {
namespace {
// The same URL / filename / SHA-256 triples COLMAP carries in
// src/colmap/feature/resources.h. Keep them identical: the point of fetching
// from COLMAP's release is that a parity run compares implementations, not
// checkpoints.
const ModelSource kSources[] = {
{"aliked-n16rot", "aliked-n16rot.onnx",
"https://github.com/colmap/colmap/releases/download/3.13.0/aliked-n16rot.onnx",
"39c423d0a6f03d39ec89d3d1d61853765c2fb6a8b8381376c703e5758778a547", 2997054ull},
{"aliked-n32", "aliked-n32.onnx",
"https://github.com/colmap/colmap/releases/download/3.13.0/aliked-n32.onnx",
"a077728a02d2de1a775c66df6de8cfeb7c6b51ca57572c64c680131c988c8b3c", 4205634ull},
{"aliked-lightglue", "aliked-lightglue.onnx",
"https://github.com/colmap/colmap/releases/download/3.13.0/aliked-lightglue.onnx",
"b9a5de7204648b18a8cf5dcac819f9d30de1a5961ef03756803c8b86c2dceb8d", 0ull},
};
// ---------------------------------------------------------------------------
// SHA-256 (FIPS 180-4)
//
// Here because the repository has no crypto dependency and should not grow one
// for this. Verifying matters more than usual: these bytes come off the
// network and go straight into a parser.
// ---------------------------------------------------------------------------
struct Sha256 {
uint32_t h[8] = {0x6a09e667u, 0xbb67ae85u, 0x3c6ef372u, 0xa54ff53au,
0x510e527fu, 0x9b05688cu, 0x1f83d9abu, 0x5be0cd19u};
uint8_t buf[64] = {};
size_t buf_len = 0;
uint64_t total = 0;
static uint32_t ror(uint32_t x, int n) { return (x >> n) | (x << (32 - n)); }
void block(const uint8_t* p) {
static const uint32_t K[64] = {
0x428a2f98u, 0x71374491u, 0xb5c0fbcfu, 0xe9b5dba5u, 0x3956c25bu, 0x59f111f1u,
0x923f82a4u, 0xab1c5ed5u, 0xd807aa98u, 0x12835b01u, 0x243185beu, 0x550c7dc3u,
0x72be5d74u, 0x80deb1feu, 0x9bdc06a7u, 0xc19bf174u, 0xe49b69c1u, 0xefbe4786u,
0x0fc19dc6u, 0x240ca1ccu, 0x2de92c6fu, 0x4a7484aau, 0x5cb0a9dcu, 0x76f988dau,
0x983e5152u, 0xa831c66du, 0xb00327c8u, 0xbf597fc7u, 0xc6e00bf3u, 0xd5a79147u,
0x06ca6351u, 0x14292967u, 0x27b70a85u, 0x2e1b2138u, 0x4d2c6dfcu, 0x53380d13u,
0x650a7354u, 0x766a0abbu, 0x81c2c92eu, 0x92722c85u, 0xa2bfe8a1u, 0xa81a664bu,
0xc24b8b70u, 0xc76c51a3u, 0xd192e819u, 0xd6990624u, 0xf40e3585u, 0x106aa070u,
0x19a4c116u, 0x1e376c08u, 0x2748774cu, 0x34b0bcb5u, 0x391c0cb3u, 0x4ed8aa4au,
0x5b9cca4fu, 0x682e6ff3u, 0x748f82eeu, 0x78a5636fu, 0x84c87814u, 0x8cc70208u,
0x90befffau, 0xa4506cebu, 0xbef9a3f7u, 0xc67178f2u};
uint32_t w[64];
for (int i = 0; i < 16; ++i)
w[i] = ((uint32_t)p[i * 4] << 24) | ((uint32_t)p[i * 4 + 1] << 16) |
((uint32_t)p[i * 4 + 2] << 8) | (uint32_t)p[i * 4 + 3];
for (int i = 16; i < 64; ++i) {
const uint32_t s0 = ror(w[i - 15], 7) ^ ror(w[i - 15], 18) ^ (w[i - 15] >> 3);
const uint32_t s1 = ror(w[i - 2], 17) ^ ror(w[i - 2], 19) ^ (w[i - 2] >> 10);
w[i] = w[i - 16] + s0 + w[i - 7] + s1;
}
uint32_t a = h[0], b = h[1], c = h[2], d = h[3];
uint32_t e = h[4], f = h[5], g = h[6], hh = h[7];
for (int i = 0; i < 64; ++i) {
const uint32_t S1 = ror(e, 6) ^ ror(e, 11) ^ ror(e, 25);
const uint32_t ch = (e & f) ^ (~e & g);
const uint32_t t1 = hh + S1 + ch + K[i] + w[i];
const uint32_t S0 = ror(a, 2) ^ ror(a, 13) ^ ror(a, 22);
const uint32_t maj = (a & b) ^ (a & c) ^ (b & c);
const uint32_t t2 = S0 + maj;
hh = g; g = f; f = e; e = d + t1;
d = c; c = b; b = a; a = t1 + t2;
}
h[0] += a; h[1] += b; h[2] += c; h[3] += d;
h[4] += e; h[5] += f; h[6] += g; h[7] += hh;
}
void update(const uint8_t* p, size_t n) {
total += n;
while (n) {
const size_t take = std::min(n, sizeof(buf) - buf_len);
std::memcpy(buf + buf_len, p, take);
buf_len += take;
p += take;
n -= take;
if (buf_len == sizeof(buf)) { block(buf); buf_len = 0; }
}
}
std::string hex() {
const uint64_t bits = total * 8;
uint8_t pad = 0x80;
update(&pad, 1);
pad = 0x00;
while (buf_len != 56) update(&pad, 1);
uint8_t len[8];
for (int i = 0; i < 8; ++i) len[i] = (uint8_t)(bits >> (56 - 8 * i));
// update() would count these into `total`, but the length is already
// frozen in `bits`, so the extra count is harmless.
update(len, 8);
static const char* kHex = "0123456789abcdef";
std::string out(64, '0');
for (int i = 0; i < 8; ++i)
for (int b = 0; b < 4; ++b) {
const uint8_t v = (uint8_t)(h[i] >> (24 - 8 * b));
out[(size_t)(i * 8 + b * 2)] = kHex[v >> 4];
out[(size_t)(i * 8 + b * 2 + 1)] = kHex[v & 0xF];
}
return out;
}
};
std::string env_str(const char* name) {
const char* v = std::getenv(name);
return v ? std::string(v) : std::string();
}
fs::path cache_root() {
#ifdef _WIN32
fs::path dir = env_str("LOCALAPPDATA");
#else
fs::path dir = env_str("XDG_CACHE_HOME");
if (dir.empty()) {
const std::string home = env_str("HOME");
if (!home.empty()) dir = fs::path(home) / ".cache";
}
#endif
if (dir.empty()) dir = ".";
return dir / "spirulae-splat";
}
bool have_curl() {
#ifdef _WIN32
return std::system("curl --version >NUL 2>&1") == 0;
#else
return std::system("curl --version >/dev/null 2>&1") == 0;
#endif
}
} // namespace
const ModelSource* find_model_source(const std::string& id) {
for (const ModelSource& s : kSources)
if (id == s.id) return &s;
return nullptr;
}
std::string model_cache_path(const ModelSource& src) {
return (cache_root() / "models" / src.file).string();
}
std::string sha256_file(const std::string& path) {
std::ifstream fin(path, std::ios::binary);
if (!fin) return {};
Sha256 sha;
std::vector<char> chunk(1 << 20);
while (fin) {
fin.read(chunk.data(), (std::streamsize)chunk.size());
const std::streamsize got = fin.gcount();
if (got > 0) sha.update(reinterpret_cast<const uint8_t*>(chunk.data()), (size_t)got);
}
return sha.hex();
}
std::string ensure_model(const ModelSource& src) {
const fs::path dst = model_cache_path(src);
std::error_code ec;
if (fs::exists(dst, ec)) {
const std::string got = sha256_file(dst.string());
if (got == src.sha256) return dst.string();
// A cached file that does not hash is a failed or interrupted download
// from a previous run, not a reason to stop: say so and refetch.
NN_LOG_WARN("[aliked] cached %s does not match its checksum; re-downloading\n",
src.file);
fs::remove(dst, ec);
}
fs::create_directories(dst.parent_path(), ec);
NN_CHECK(!ec, "cannot create %s: %s", dst.parent_path().string().c_str(),
ec.message().c_str());
NN_CHECK(have_curl(),
"curl was not found, and it is how checkpoints are fetched.\n"
" Install curl, or download\n %s\n to\n %s\n by hand.",
src.url, dst.string().c_str());
fs::path part = dst;
part += ".part";
NN_LOG_INFO("[aliked] fetching %s (%.1f MB) from %s\n", src.file,
(double)src.bytes / 1e6, src.url);
// -C - resumes a partial .part file; -f makes an HTTP error an exit code
// rather than a saved error page. Downloading into .part and renaming is
// what keeps an interrupted fetch from ever looking like a complete model.
std::string cmd = "curl -L -f --progress-bar -C - -o \"" + part.string() + "\" \"" +
std::string(src.url) + "\"";
const int rc = std::system(cmd.c_str());
if (rc != 0) {
fs::remove(part, ec);
nn::fail("downloading %s failed (curl exit %d).\n"
" Fetch it by hand from\n %s\n and save it as\n %s",
src.file, rc, src.url, dst.string().c_str());
}
const std::string got = sha256_file(part.string());
if (got != src.sha256) {
fs::remove(part, ec);
nn::fail("%s downloaded but its SHA-256 is\n %s\n expected\n %s\n"
" The file was discarded.",
src.file, got.c_str(), src.sha256);
}
fs::rename(part, dst, ec);
NN_CHECK(!ec, "cannot move the download into place: %s", ec.message().c_str());
NN_LOG_INFO("[aliked] saved %s\n", dst.string().c_str());
return dst.string();
}
std::string resolve_model(const std::string& id_or_path) {
if (const ModelSource* src = find_model_source(id_or_path)) return ensure_model(*src);
std::error_code ec;
NN_CHECK(fs::exists(id_or_path, ec),
"'%s' is neither a known model id (aliked-n16rot, aliked-n32, "
"aliked-lightglue) nor a file that exists",
id_or_path.c_str());
return id_or_path;
}
} // namespace aliked
src/aliked/model/Fetch.h
+55
View File
@@ -0,0 +1,55 @@
#pragma once
// Getting a checkpoint onto disk.
//
// The artifacts are COLMAP's, fetched from COLMAP's release URLs and verified
// against COLMAP's SHA-256 (docs/notes/aliked-port-plan.md). Nothing is
// bundled, nothing is committed, and nothing of ours is hosted anywhere -- so
// there is no converter to keep in step with an upstream re-export, and the
// bytes we run are the bytes the reference implementation runs.
//
// Both models are permissively licensed -- ALIKED is BSD-3-Clause, LightGlue
// Apache-2.0 -- so unlike the segmentation checkpoints
// (src/app/gui/ModelCache.cpp) this needs no consent gate. It still never
// downloads behind the user's back: `ssplat sfm` prints what it is fetching
// and from where, and --aliked-model points at a file instead.
#include <cstdint>
#include <string>
namespace aliked {
// A checkpoint we know how to fetch. `sha256` is lowercase hex, and is checked
// after download and on every subsequent load of the cached file -- a
// truncated or tampered artifact must not reach the parser.
struct ModelSource {
const char* id; // "aliked-n16rot" -- what --features spells
const char* file; // basename in the cache directory
const char* url;
const char* sha256;
uint64_t bytes; // approximate, for the "downloading N MB" line
};
// Null when `id` is not one of ours.
const ModelSource* find_model_source(const std::string& id);
// Where a cached checkpoint lives: <cache>/spirulae-splat/models/<file>.
// Mirrors src/app/gui/AppPaths.cpp's cache_dir(); duplicated rather than
// shared because src/aliked/ sits below src/app/ in the layering and may not
// include it.
std::string model_cache_path(const ModelSource& src);
// Returns a path to a verified local copy, downloading through the system
// `curl` if needed. Throws nn::Error with an actionable message -- including
// the URL to fetch by hand -- when curl is missing, the download fails, or the
// hash does not match.
std::string ensure_model(const ModelSource& src);
// Resolve what the user asked for: an explicit path is used as-is (and, when
// it happens to be one of the known files, still hash-checked), otherwise the
// id is fetched.
std::string resolve_model(const std::string& id_or_path);
// Lowercase hex SHA-256 of a file's contents. Empty when it cannot be read.
std::string sha256_file(const std::string& path);
} // namespace aliked
src/aliked/model/LightGlue.cpp
+537
View File
@@ -0,0 +1,537 @@
// LightGlue's forward pass, next to its PyTorch reference (cvg/LightGlue,
// lightglue.py). Nine layers of self- then cross-attention, one assignment.
//
// Three things about the checkpoint decide how this file is written:
//
// * torch exports nn.Linear as MatMul + Add, and only the Add's bias keeps
// its qualified name -- the weight arrives as "onnx::MatMul_2537", and its
// matrix is [in, out] rather than the [out, in] every op here wants. Both
// are fixed once, at load: OnnxFile::linearWeights() recovers the pairing
// by walking Add -> MatMul, and the matrices are transposed on the host.
// * the fused qkv projection is laid out [head][dim][3] -- q, k and v
// INTERLEAVED per element, which no stride can express. Permuting the
// projection's output rows at load makes it [3][head][dim] instead, which
// is exactly the fused layout nn::attention's q/k/v strides address.
// * the export runs the assignment head on the last layer only. LightGlue's
// early exit and token pruning are not in the graph, so they are not here
// either; both are speedups, not behaviour.
#include "aliked/model/LightGlue.h"
#include "aliked/Common.h"
#include "aliked/model/Fetch.h"
#include "aliked/model/Onnx.h"
#include "nn/Ops.h"
#include "nn/Tensor.h"
#include "nn/vk/EmbeddedSpirv.h"
#include "nn/vk/Memory.h"
#include "nn/vk/Stream.h"
#include <algorithm>
#include <cmath>
#include <cstdio>
#include <cstring>
#include <unordered_map>
NN_DECLARE_EMBEDDED_MODULES(aliked)
namespace aliked {
namespace {
using nn::Act;
using nn::AttnOpts;
using nn::DType;
using nn::LinearOpts;
using nn::Tensor;
constexpr int64_t kAlign = 256;
int64_t align_up(int64_t v, int64_t a) { return (v + a - 1) / a * a; }
struct AssignParams {
uint64_t out_val, out_idx, sim, bias;
uint32_t R, C, stride, step;
float scale;
uint32_t groups_per_row;
};
} // namespace
struct Matcher::Impl {
// ---- hyperparameters, read off the checkpoint ----
int n_layers = 0;
int dim = 0; // 256
int input_dim = 0; // 128, ALIKED's descriptor width
int n_heads = 4; // not stated in the file; LightGlue's every release
int head_dim = 0;
std::unordered_map<std::string, Tensor> w;
nn::DevicePtr blob = 0;
vk::Arena arena{"lightglue"};
std::string path;
bool loaded = false;
~Impl() {
if (blob) vk::device_free(blob);
}
Tensor get(const std::string& name) const {
auto it = w.find(name);
NN_CHECK(it != w.end(), "LightGlue checkpoint has no '%s'", name.c_str());
return it->second;
}
Tensor getf(const char* fmt, int i, const char* suffix) const {
char buf[128];
std::snprintf(buf, sizeof buf, fmt, i, suffix);
return get(buf);
}
void load(const std::string& onnx_path);
std::vector<Match> match(const MatchInput& a, const MatchInput& b,
const MatchOptions& opts);
// x <- x + ffn(cat(x, message)); the FFN is Linear, LayerNorm, GELU, Linear.
void ffn(const Tensor& x, const Tensor& message, int64_t n, const char* prefix,
int layer);
};
// ---------------------------------------------------------------------------
// Loading
// ---------------------------------------------------------------------------
void Matcher::Impl::load(const std::string& onnx_path) {
NN_CHECK(!loaded, "LightGlue::load called twice");
const OnnxFile file = read_onnx(onnx_path);
const auto linears = file.linearWeights();
const char* path = onnx_path.c_str();
// How many layers the file has, from the highest transformers.N present.
n_layers = 0;
for (const auto& kv : linears) {
int idx = -1;
if (std::sscanf(kv.first.c_str(), "transformers.%d.", &idx) == 1 && idx >= 0)
n_layers = std::max(n_layers, idx + 1);
}
NN_CHECK(n_layers > 0, "'%s' has no transformer layers; is this a LightGlue "
"checkpoint?", path);
// Staged host copies, transposed / permuted as described at the top.
struct Staged {
std::string name;
std::vector<int64_t> shape;
std::vector<float> data;
};
std::vector<Staged> staged;
auto init_of = [&](const std::string& n) -> const OnnxTensor& {
const OnnxTensor* t = file.find(n);
NN_CHECK(t != nullptr, "'%s' has no initializer '%s'", path, n.c_str());
return *t;
};
// A Linear: weight [in, out] in the file -> [out, in] here, plus its bias.
auto take_linear = [&](const std::string& module) {
auto it = linears.find(module);
NN_CHECK(it != linears.end(), "'%s' has no linear named '%s'", path,
module.c_str());
const OnnxTensor& wt = init_of(it->second);
NN_CHECK(wt.shape.size() == 2, "'%s': %s is not a matrix", path, module.c_str());
const int64_t in = wt.shape[0], out = wt.shape[1];
Staged s{module + ".weight", {out, in}, std::vector<float>(wt.data.size())};
for (int64_t i = 0; i < in; ++i)
for (int64_t o = 0; o < out; ++o)
s.data[(size_t)(o * in + i)] = wt.data[(size_t)(i * out + o)];
staged.push_back(std::move(s));
const OnnxTensor& b = init_of(module + ".bias");
staged.push_back(Staged{module + ".bias", b.shape, b.data});
};
auto take_raw = [&](const std::string& name) {
const OnnxTensor& t = init_of(name);
staged.push_back(Staged{name, t.shape, t.data});
};
take_linear("input_proj");
input_dim = (int)staged[staged.size() - 2].shape[1];
dim = (int)staged[staged.size() - 2].shape[0];
head_dim = dim / n_heads;
NN_CHECK(head_dim * n_heads == dim && head_dim % 2 == 0,
"'%s': %d channels do not split into %d rotary heads", path, dim, n_heads);
for (int l = 0; l < n_layers; ++l) {
char p[96];
std::snprintf(p, sizeof p, "transformers.%d.self_attn", l);
const std::string sa = p;
// Wqkv, with its output rows permuted from [head][dim][3] to
// [3][head][dim] so q, k and v become contiguous blocks.
{
auto it = linears.find(sa + ".Wqkv");
NN_CHECK(it != linears.end(), "'%s' has no %s.Wqkv", path, sa.c_str());
const OnnxTensor& wt = init_of(it->second);
const OnnxTensor& bt = init_of(sa + ".Wqkv.bias");
const int64_t in = wt.shape[0], out = wt.shape[1];
NN_CHECK(out == 3 * dim, "'%s': %s.Wqkv is %lld wide, expected %lld", path,
sa.c_str(), (long long)out, (long long)(3 * dim));
Staged sw{sa + ".Wqkv.weight", {out, in}, std::vector<float>(wt.data.size())};
Staged sb{sa + ".Wqkv.bias", {out}, std::vector<float>((size_t)out)};
for (int h = 0; h < n_heads; ++h)
for (int d = 0; d < head_dim; ++d)
for (int j = 0; j < 3; ++j) {
const int64_t src = (int64_t)(h * head_dim + d) * 3 + j;
const int64_t dst = (int64_t)j * dim + h * head_dim + d;
for (int64_t i = 0; i < in; ++i)
sw.data[(size_t)(dst * in + i)] =
wt.data[(size_t)(i * out + src)];
sb.data[(size_t)dst] = bt.data[(size_t)src];
}
staged.push_back(std::move(sw));
staged.push_back(std::move(sb));
}
take_linear(sa + ".out_proj");
take_linear(sa + ".ffn.0");
take_raw(sa + ".ffn.1.weight");
take_raw(sa + ".ffn.1.bias");
take_linear(sa + ".ffn.3");
std::snprintf(p, sizeof p, "transformers.%d.cross_attn", l);
const std::string ca = p;
take_linear(ca + ".to_qk");
take_linear(ca + ".to_v");
take_linear(ca + ".to_out");
take_linear(ca + ".ffn.0");
take_raw(ca + ".ffn.1.weight");
take_raw(ca + ".ffn.1.bias");
take_linear(ca + ".ffn.3");
}
{
char p[64];
std::snprintf(p, sizeof p, "log_assignment.%d", n_layers - 1);
take_linear(std::string(p) + ".final_proj");
take_linear(std::string(p) + ".matchability");
}
// The Fourier positional encoding's Wr is a bias-free Linear(2, dim/2), so
// the Add-walk above cannot find it. It is the only [2, N] initializer in
// the file, which is enough to name it.
{
const OnnxTensor* wr = nullptr;
for (const OnnxTensor& t : file.initializers)
if (t.shape.size() == 2 && t.shape[0] == 2) {
NN_CHECK(wr == nullptr,
"'%s' has more than one [2, N] initializer; the positional "
"encoding can no longer be identified by shape",
path);
wr = &t;
}
NN_CHECK(wr != nullptr, "'%s' has no positional-encoding matrix", path);
NN_CHECK(wr->shape[1] * 2 == head_dim,
"'%s': positional encoding is [2, %lld], expected [2, %d]", path,
(long long)wr->shape[1], head_dim / 2);
Staged s{"posenc.weight", {wr->shape[1], 2}, std::vector<float>(wr->data.size())};
for (int64_t i = 0; i < 2; ++i)
for (int64_t o = 0; o < wr->shape[1]; ++o)
s.data[(size_t)(o * 2 + i)] = wr->data[(size_t)(i * wr->shape[1] + o)];
staged.push_back(std::move(s));
}
int64_t total = 0;
for (const Staged& s : staged) total = align_up(total, kAlign) + (int64_t)s.data.size() * 4;
blob = vk::device_alloc((uint64_t)total, "lightglue-weights");
int64_t off = 0;
for (const Staged& s : staged) {
off = align_up(off, kAlign);
const nn::DevicePtr ptr = blob + (uint64_t)off;
vk::Stream::get().upload(ptr, s.data.data(), s.data.size() * 4);
off += (int64_t)s.data.size() * 4;
Tensor t;
t.ptr = ptr;
t.dtype = DType::F32;
t.ndim = (int32_t)s.shape.size();
for (size_t i = 0; i < s.shape.size() && i < 4; ++i) t.shape[i] = s.shape[i];
w[s.name] = t;
}
vk::Stream::get().sync();
this->path = onnx_path;
loaded = true;
NN_LOG_INFO("[lightglue] %s: %d layers, dim %d (%d heads), input %d, %zu tensors, "
"%.1f MB on device\n",
onnx_path.c_str(), n_layers, dim, n_heads, input_dim, w.size(),
(double)total / 1e6);
}
// ---------------------------------------------------------------------------
// Forward
// ---------------------------------------------------------------------------
void Matcher::Impl::ffn(const Tensor& x, const Tensor& message, int64_t n,
const char* prefix, int layer) {
char p[128];
std::snprintf(p, sizeof p, "transformers.%d.%s", layer, prefix);
const std::string b = p;
vk::ArenaScope scope(arena);
// cat(x, message) along the channel axis, then Linear(2d, 2d).
Tensor cat = nn::arena_tensor(arena, DType::F32, n, 2 * dim);
nn::strided_copy(cat, x, n, dim, dim, 2 * dim);
nn::strided_copy(cat.offsetElems(dim), message, n, dim, dim, 2 * dim);
Tensor h = nn::arena_tensor(arena, DType::F32, n, 2 * dim);
LinearOpts l0;
l0.bias = get(b + ".ffn.0.bias");
nn::linear(h, cat, get(b + ".ffn.0.weight"), l0);
nn::layer_norm(h, h, get(b + ".ffn.1.weight"), get(b + ".ffn.1.bias"), 1e-5f);
nn::unary(h, h, Act::GeluErf);
// The residual is x itself, so the final projection accumulates into it.
LinearOpts l1;
l1.bias = get(b + ".ffn.3.bias");
l1.residual = x;
nn::linear(x, h, get(b + ".ffn.3.weight"), l1);
}
std::vector<Match> Matcher::Impl::match(const MatchInput& A, const MatchInput& B,
const MatchOptions& opts) {
std::vector<Match> out;
const int64_t n0 = A.n, n1 = B.n;
if (n0 == 0 || n1 == 0) return out;
// Peak live set: two [n, dim] states, one [n, 3*dim] projection, the
// similarity matrix, and the FFN's [n, 2*dim] scratch.
const int64_t nmax = std::max(n0, n1);
arena.reserve((uint64_t)((n0 * n1 + nmax * (10 * dim + 8)) * 4 + (32 << 20)));
vk::ArenaScope root(arena);
Tensor x[2];
Tensor freqs[2];
const MatchInput* in[2] = {&A, &B};
const int64_t nn_[2] = {n0, n1};
for (int s = 0; s < 2; ++s) {
const MatchInput& I = *in[s];
const int64_t n = nn_[s];
// Keypoints: our corner origin -> LightGlue's centre origin, then its
// own normalization -- shift by the image centre, scale by half the
// LONGER side (not per-axis, so the aspect ratio survives).
std::vector<float> kn((size_t)n * 2);
const float sx = 0.5f * (float)I.width, sy = 0.5f * (float)I.height;
const float scale = 0.5f * (float)std::max(I.width, I.height);
for (int64_t i = 0; i < n; ++i) {
kn[(size_t)i * 2] = (I.keypoints[i * 2] - 0.5f - sx) / scale;
kn[(size_t)i * 2 + 1] = (I.keypoints[i * 2 + 1] - 0.5f - sy) / scale;
}
Tensor tk = nn::arena_tensor(arena, DType::F32, n, 2);
nn::tensor_from_host(tk, kn.data(), (int64_t)kn.size());
// Rotary table: proj = kn @ Wr^T, then (cos, sin) per pair, shared by
// every head -- which is exactly nn::rope's [n, head_dim/2, 2] layout.
Tensor proj = nn::arena_tensor(arena, DType::F32, n, head_dim / 2);
nn::linear(proj, tk, get("posenc.weight"));
freqs[s] = nn::arena_tensor(arena, DType::F32, n * (head_dim / 2), 2);
{
std::vector<float> host((size_t)n * (head_dim / 2));
nn::tensor_to_host(proj, host.data(), (int64_t)host.size());
std::vector<float> cs(host.size() * 2);
for (size_t i = 0; i < host.size(); ++i) {
cs[i * 2] = std::cos(host[i]);
cs[i * 2 + 1] = std::sin(host[i]);
}
nn::tensor_from_host(freqs[s], cs.data(), (int64_t)cs.size());
}
Tensor desc = nn::arena_tensor(arena, DType::F32, n, input_dim);
nn::tensor_from_host(desc, I.descriptors, n * input_dim);
x[s] = nn::arena_tensor(arena, DType::F32, n, dim);
LinearOpts lp;
lp.bias = get("input_proj.bias");
nn::linear(x[s], desc, get("input_proj.weight"), lp);
nn::tensor_debug_dump(s ? "lg input_proj[1]" : "lg input_proj[0]", x[s]);
}
for (int l = 0; l < n_layers; ++l) {
// ---- self-attention, each image independently ----
for (int s = 0; s < 2; ++s) {
const int64_t n = nn_[s];
vk::ArenaScope scope(arena);
char p[96];
std::snprintf(p, sizeof p, "transformers.%d.self_attn", l);
const std::string sa = p;
Tensor qkv = nn::arena_tensor(arena, DType::F32, n, 3 * dim);
LinearOpts lq;
lq.bias = get(sa + ".Wqkv.bias");
nn::linear(qkv, x[s], get(sa + ".Wqkv.weight"), lq);
// q and k rotate in place inside the fused buffer; v does not.
nn::rope(qkv, freqs[s], n_heads, head_dim, n, 1, 3 * dim);
nn::rope(qkv.offsetElems(dim), freqs[s], n_heads, head_dim, n, 1, 3 * dim);
Tensor ctx = nn::arena_tensor(arena, DType::F32, n, dim);
AttnOpts ao;
ao.n_heads = n_heads;
ao.head_dim = head_dim;
ao.q_stride = ao.k_stride = ao.v_stride = 3 * dim;
ao.arena = &arena;
nn::attention(ctx, qkv, qkv.offsetElems(dim), qkv.offsetElems(2 * dim), n, n,
ao);
Tensor msg = nn::arena_tensor(arena, DType::F32, n, dim);
LinearOpts lo;
lo.bias = get(sa + ".out_proj.bias");
nn::linear(msg, ctx, get(sa + ".out_proj.weight"), lo);
ffn(x[s], msg, n, "self_attn", l);
}
// ---- cross-attention, both directions off the same projections ----
{
vk::ArenaScope scope(arena);
char p[96];
std::snprintf(p, sizeof p, "transformers.%d.cross_attn", l);
const std::string ca = p;
Tensor qk[2], v[2], msg[2];
for (int s = 0; s < 2; ++s) {
qk[s] = nn::arena_tensor(arena, DType::F32, nn_[s], dim);
v[s] = nn::arena_tensor(arena, DType::F32, nn_[s], dim);
LinearOpts lqk, lv;
lqk.bias = get(ca + ".to_qk.bias");
lv.bias = get(ca + ".to_v.bias");
nn::linear(qk[s], x[s], get(ca + ".to_qk.weight"), lqk);
nn::linear(v[s], x[s], get(ca + ".to_v.weight"), lv);
}
for (int s = 0; s < 2; ++s) {
const int o = 1 - s;
Tensor ctx = nn::arena_tensor(arena, DType::F32, nn_[s], dim);
AttnOpts ao;
ao.n_heads = n_heads;
ao.head_dim = head_dim;
ao.arena = &arena;
// No rotary here: the cross block attends between images, where
// a within-image position has no meaning.
nn::attention(ctx, qk[s], qk[o], v[o], nn_[s], nn_[o], ao);
msg[s] = nn::arena_tensor(arena, DType::F32, nn_[s], dim);
LinearOpts lo;
lo.bias = get(ca + ".to_out.bias");
nn::linear(msg[s], ctx, get(ca + ".to_out.weight"), lo);
}
for (int s = 0; s < 2; ++s) ffn(x[s], msg[s], nn_[s], "cross_attn", l);
}
if (l == 0 || l == n_layers - 1) {
char lbl[48];
std::snprintf(lbl, sizeof lbl, "lg layer%d[0]", l);
nn::tensor_debug_dump(lbl, x[0]);
}
}
// ---- assignment ----
char pfx[64];
std::snprintf(pfx, sizeof pfx, "log_assignment.%d", n_layers - 1);
const std::string la = pfx;
Tensor md[2], z[2];
for (int s = 0; s < 2; ++s) {
md[s] = nn::arena_tensor(arena, DType::F32, nn_[s], dim);
LinearOpts lf;
lf.bias = get(la + ".final_proj.bias");
nn::linear(md[s], x[s], get(la + ".final_proj.weight"), lf);
z[s] = nn::arena_tensor(arena, DType::F32, nn_[s], 1, 1, 1, 2);
LinearOpts lm;
lm.bias = get(la + ".matchability.bias");
lm.act = Act::LogSigmoid; // logsigmoid(z), the certainty term
nn::linear(z[s], x[s], get(la + ".matchability.weight"), lm);
}
// sim = (md0 / d^0.25) . (md1 / d^0.25), i.e. one 1/sqrt(dim) on the product.
Tensor sim = nn::arena_tensor(arena, DType::F32, n0, n1);
nn::matmul_nt(sim, md[0], md[1], 1.0f / std::sqrt((float)dim));
nn::tensor_debug_dump("lg mdesc0", md[0]);
nn::tensor_debug_dump("lg z0", z[0].view(n0));
nn::tensor_debug_dump("lg sim", sim);
// The row and column log-sum-exps, and then the two arg-maxes with every
// column-dependent term folded into a bias. See shaders/aliked.slang for
// why the score matrix is never written out.
Tensor lse0 = nn::arena_tensor(arena, DType::F32, n0);
Tensor lse1 = nn::arena_tensor(arena, DType::F32, n1);
Tensor bias0 = nn::arena_tensor(arena, DType::F32, n1);
Tensor bias1 = nn::arena_tensor(arena, DType::F32, n0);
Tensor val0 = nn::arena_tensor(arena, DType::F32, n0);
Tensor val1 = nn::arena_tensor(arena, DType::F32, n1);
Tensor idx0 = nn::arena_tensor(arena, DType::I32, n0);
Tensor idx1 = nn::arena_tensor(arena, DType::I32, n1);
auto run = [&](const char* entry, const Tensor& ov, const Tensor& oi,
const Tensor& bias, int64_t R, int64_t C, int64_t stride,
int64_t step, float scale = 1.0f) {
AssignParams p{};
p.out_val = ov.ptr;
p.out_idx = oi.valid() ? oi.ptr : vk::or_fallback(0);
p.sim = sim.ptr;
p.bias = bias.valid() ? bias.ptr : vk::or_fallback(0);
p.R = (uint32_t)R;
p.C = (uint32_t)C;
p.stride = (uint32_t)stride;
p.step = (uint32_t)step;
p.scale = scale;
const vk::Stream::Fold fold = vk::Stream::fold1D(R, 1);
p.groups_per_row = fold.per_row;
vk::Stream::get().dispatch(entry, vk::SpecList{0u, 0u}, fold.per_row, fold.rows, 1,
&p, sizeof(p));
};
run("aliked.assign_logsumexp", lse0, {}, {}, n0, n1, n1, 1);
run("aliked.assign_logsumexp", lse1, {}, {}, n1, n0, 1, n1);
// bias for a row scan is (logsigmoid(z1) - lse1) per column, and vice versa.
nn::add(bias0, z[1].view(n1), lse1, 1.0f, -1.0f);
nn::add(bias1, z[0].view(n0), lse0, 1.0f, -1.0f);
// 2 * sim, because the score sums a row log-softmax and a column one and
// each contains sim. See the kernel.
run("aliked.assign_argmax", val0, idx0, bias0, n0, n1, n1, 1, 2.0f);
run("aliked.assign_argmax", val1, idx1, bias1, n1, n0, 1, n1, 2.0f);
nn::tensor_debug_dump("lg lse0", lse0);
nn::tensor_debug_dump("lg lse1", lse1);
nn::tensor_debug_dump("lg val0", val0);
std::vector<float> h_val0((size_t)n0), h_lse0((size_t)n0), h_z0((size_t)n0);
std::vector<int32_t> h_idx0((size_t)n0), h_idx1((size_t)n1);
nn::tensor_to_host(val0, h_val0.data(), n0);
nn::tensor_to_host(lse0, h_lse0.data(), n0);
nn::tensor_to_host(z[0].view(n0), h_z0.data(), n0);
vk::Stream::get().download(h_idx0.data(), idx0.ptr, (uint64_t)n0 * 4);
vk::Stream::get().download(h_idx1.data(), idx1.ptr, (uint64_t)n1 * 4);
out.reserve((size_t)std::min(n0, n1));
for (int64_t i = 0; i < n0; ++i) {
const int32_t j = h_idx0[(size_t)i];
if (j < 0 || j >= n1) continue;
if (h_idx1[(size_t)j] != (int32_t)i) continue; // mutual nearest
// The row-constant terms, which the arg-max did not need.
const float score = std::exp(h_val0[(size_t)i] - h_lse0[(size_t)i] +
h_z0[(size_t)i]);
if (score < opts.min_score) continue;
out.push_back({(uint32_t)i, (uint32_t)j, score});
}
return out;
}
// ---------------------------------------------------------------------------
Matcher::Matcher() : impl_(new Impl) {}
Matcher::~Matcher() { delete impl_; }
bool Matcher::loaded() const { return impl_->loaded; }
int Matcher::descriptorDim() const { return impl_->input_dim; }
void Matcher::load(const std::string& model) {
NN_ENSURE_EMBEDDED_MODULES(aliked);
impl_->load(resolve_model(model));
}
std::vector<Match> Matcher::match(const MatchInput& a, const MatchInput& b,
const MatchOptions& opts) {
NN_CHECK(impl_->loaded, "LightGlue::match before load()");
NN_CHECK((int)a.n == 0 || impl_->input_dim > 0, "LightGlue: no descriptor width");
return impl_->match(a, b, opts);
}
} // namespace aliked
src/aliked/model/LightGlue.h
+65
View File
@@ -0,0 +1,65 @@
#pragma once
// LightGlue: a learned matcher for ALIKED keypoints.
//
// Two sets of keypoints and descriptors in, a list of correspondences out. The
// network is nine transformer layers, each doing self-attention within each
// image and cross-attention between them, followed by one assignment head.
//
// It needed NO new general ops -- linear, layer_norm, attention and rope cover
// the whole forward pass, and our RoPE frequency layout is already the one
// LightGlue's rotary encoding uses. Only the assignment tail is here, in
// shaders/aliked.slang, and only because it is a reduction over a matrix that
// should never be materialized.
//
// Cost is the thing to know before using it: roughly 100 GFLOP for a
// 2048 x 2048 pair, i.e. tens of milliseconds. It is a matcher for a
// SHORTLIST, never for a raw exhaustive pair list -- run pair selection first.
#include <cstdint>
#include <string>
#include <vector>
namespace aliked {
struct MatchOptions {
// Assignment confidence a pair must reach. COLMAP's
// LightGlueONNXMatchingOptions default.
float min_score = 0.1f;
};
struct Match {
uint32_t i = 0, j = 0; // indices into set 0 and set 1
float score = 0;
};
// One image's input. Keypoints are in pixels with the top-left pixel's CORNER
// at (0, 0) -- this repository's convention, which is half a pixel off
// LightGlue's own; the conversion happens inside.
struct MatchInput {
const float* keypoints = nullptr; // [n, 2] xy
const float* descriptors = nullptr; // [n, dim], L2-normalized
uint32_t n = 0;
int width = 0, height = 0; // the image the keypoints refer to
};
class Matcher {
public:
Matcher();
~Matcher();
Matcher(const Matcher&) = delete;
Matcher& operator=(const Matcher&) = delete;
// "aliked-lightglue" (fetched and cached) or a path to an .onnx file.
void load(const std::string& model);
bool loaded() const;
int descriptorDim() const;
std::vector<Match> match(const MatchInput& a, const MatchInput& b,
const MatchOptions& opts = {});
private:
struct Impl;
Impl* impl_ = nullptr;
};
} // namespace aliked
src/aliked/model/Onnx.cpp
+374
View File
@@ -0,0 +1,374 @@
#include "aliked/model/Onnx.h"
#include "nn/core/Error.h"
#include "nn/core/Half.h"
#include "nn/core/Log.h"
#include <cstring>
#include <fstream>
#include <unordered_set>
namespace aliked {
namespace {
// ONNX TensorProto.DataType, the handful that can appear as a weight.
enum : int32_t {
kFloat = 1,
kUint8 = 2,
kInt8 = 3,
kInt32 = 6,
kInt64 = 7,
kDouble = 11,
kFloat16 = 10,
};
const char* dtype_name(int32_t t) {
switch (t) {
case kFloat: return "float";
case kUint8: return "uint8";
case kInt8: return "int8";
case kInt32: return "int32";
case kInt64: return "int64";
case kDouble: return "double";
case kFloat16: return "float16";
default: return "?";
}
}
// A bounds-checked cursor over the file. Every read validates against `end`
// before it advances, so a truncated or hostile file throws instead of walking
// off the buffer -- this parses bytes fetched over the network.
struct Reader {
const uint8_t* p = nullptr;
const uint8_t* end = nullptr;
const char* what = "";
bool done() const { return p >= end; }
[[noreturn]] void bad(const char* why) const {
nn::fail("%s: malformed ONNX (%s)", what, why);
}
uint64_t varint() {
uint64_t r = 0;
int shift = 0;
while (true) {
if (p >= end) bad("truncated varint");
const uint8_t c = *p++;
// 10 groups of 7 bits covers a full uint64; more means corruption.
if (shift > 63) bad("varint too long");
r |= (uint64_t)(c & 0x7F) << shift;
if (!(c & 0x80)) return r;
shift += 7;
}
}
void skip(uint64_t n) {
if ((uint64_t)(end - p) < n) bad("truncated field");
p += n;
}
};
// One protobuf field: its number, and a sub-reader over its payload (wire type
// 2) or its integer value (wire types 0/1/5).
struct Field {
uint32_t number = 0;
uint32_t wire = 0;
uint64_t value = 0; // wire 0/1/5
Reader bytes; // wire 2
};
bool next_field(Reader& r, Field& f) {
if (r.done()) return false;
const uint64_t key = r.varint();
f.number = (uint32_t)(key >> 3);
f.wire = (uint32_t)(key & 7);
switch (f.wire) {
case 0:
f.value = r.varint();
return true;
case 1:
if (r.end - r.p < 8) r.bad("truncated fixed64");
std::memcpy(&f.value, r.p, 8);
r.p += 8;
return true;
case 5: {
if (r.end - r.p < 4) r.bad("truncated fixed32");
uint32_t v32 = 0;
std::memcpy(&v32, r.p, 4);
r.p += 4;
f.value = v32;
return true;
}
case 2: {
const uint64_t n = r.varint();
if ((uint64_t)(r.end - r.p) < n) r.bad("truncated length-delimited field");
f.bytes = Reader{r.p, r.p + n, r.what};
r.p += n;
return true;
}
default:
// 3/4 are the deprecated group encoding; nothing emits them.
r.bad("unsupported wire type");
}
}
std::string to_string(const Reader& r) {
return std::string(reinterpret_cast<const char*>(r.p), (size_t)(r.end - r.p));
}
float bits_to_float(uint64_t bits32) {
const uint32_t u = (uint32_t)bits32;
float f;
std::memcpy(&f, &u, 4);
return f;
}
// Convert a raw_data payload of `dt` into f32. Element counts are checked
// against the declared shape by the caller.
void decode_raw(const Reader& raw, int32_t dt, int64_t count, std::vector<float>& out,
const char* name) {
const size_t bytes = (size_t)(raw.end - raw.p);
auto need = [&](size_t elem_size) {
if (bytes != (size_t)count * elem_size)
nn::fail("initializer '%s': %zu raw bytes for %lld %s elements", name, bytes,
(long long)count, dtype_name(dt));
};
out.resize((size_t)count);
// raw_data is little-endian by specification. Every platform this builds
// for is little-endian, so these are memcpys; a big-endian port would swap
// here and nowhere else.
switch (dt) {
case kFloat: {
need(4);
std::memcpy(out.data(), raw.p, bytes);
break;
}
case kFloat16: {
need(2);
for (int64_t i = 0; i < count; ++i) {
uint16_t h;
std::memcpy(&h, raw.p + (size_t)i * 2, 2);
out[(size_t)i] = nn::half_to_float(h);
}
break;
}
case kDouble: {
need(8);
for (int64_t i = 0; i < count; ++i) {
double d;
std::memcpy(&d, raw.p + (size_t)i * 8, 8);
out[(size_t)i] = (float)d;
}
break;
}
case kInt64: {
need(8);
for (int64_t i = 0; i < count; ++i) {
int64_t v;
std::memcpy(&v, raw.p + (size_t)i * 8, 8);
out[(size_t)i] = (float)v;
}
break;
}
case kInt32: {
need(4);
for (int64_t i = 0; i < count; ++i) {
int32_t v;
std::memcpy(&v, raw.p + (size_t)i * 4, 4);
out[(size_t)i] = (float)v;
}
break;
}
default:
nn::fail("initializer '%s' has dtype %s, which this reader does not convert",
name, dtype_name(dt));
}
}
OnnxTensor read_tensor_proto(Reader r) {
OnnxTensor t;
int32_t dt = 0;
Reader raw{}, float_data{};
bool has_raw = false, has_float = false;
Field f;
while (next_field(r, f)) {
switch (f.number) {
case 1: // dims: repeated int64, packed or not
if (f.wire == 0) {
t.shape.push_back((int64_t)f.value);
} else if (f.wire == 2) {
Reader d = f.bytes;
while (!d.done()) t.shape.push_back((int64_t)d.varint());
}
break;
case 2: // data_type
dt = (int32_t)f.value;
break;
case 4: // float_data: repeated float, packed
if (f.wire == 2) { float_data = f.bytes; has_float = true; }
break;
case 8: // name
if (f.wire == 2) t.name = to_string(f.bytes);
break;
case 9: // raw_data
if (f.wire == 2) { raw = f.bytes; has_raw = true; }
break;
default:
break; // segment, string_data, external_data, doc_string, ...
}
}
const char* name = t.name.empty() ? "<unnamed>" : t.name.c_str();
const int64_t count = t.numel();
if (count < 0 || count > (int64_t)1 << 32)
nn::fail("initializer '%s' declares %lld elements", name, (long long)count);
if (has_raw) {
decode_raw(raw, dt, count, t.data, name);
} else if (has_float && dt == kFloat) {
// The non-raw encoding: 4-byte little-endian floats, packed.
const size_t bytes = (size_t)(float_data.end - float_data.p);
if (bytes != (size_t)count * 4)
nn::fail("initializer '%s': %zu float_data bytes for %lld elements", name,
bytes, (long long)count);
t.data.resize((size_t)count);
std::memcpy(t.data.data(), float_data.p, bytes);
} else {
nn::fail("initializer '%s' carries no data this reader can read (dtype %s)", name,
dtype_name(dt));
}
return t;
}
// NodeProto: input is repeated field 1, output repeated field 2, name field 3,
// op_type field 4, attribute field 5.
//
// Two things are wanted from it. The epsilon of each BatchNormalization, keyed
// by the scale input's name, so folding BN into a conv uses the number in the
// file rather than assuming PyTorch's default. And the node's shape, so an
// anonymous initializer can be traced back to the module that consumes it
// (see OnnxFile::linearWeights).
void scan_node(Reader r, OnnxNode& out, std::unordered_map<std::string, float>& bn_eps) {
std::vector<std::string> inputs;
std::string op;
float eps = 0.0f;
bool has_eps = false;
Field f;
while (next_field(r, f)) {
if (f.number == 1 && f.wire == 2) {
inputs.push_back(to_string(f.bytes));
} else if (f.number == 2 && f.wire == 2) {
out.outputs.push_back(to_string(f.bytes));
} else if (f.number == 3 && f.wire == 2) {
out.name = to_string(f.bytes);
} else if (f.number == 4 && f.wire == 2) {
op = to_string(f.bytes);
} else if (f.number == 5 && f.wire == 2) {
// AttributeProto: name field 1, f (float) field 2.
Reader a = f.bytes;
std::string aname;
float av = 0.0f;
bool has_av = false;
Field af;
while (next_field(a, af)) {
if (af.number == 1 && af.wire == 2) aname = to_string(af.bytes);
else if (af.number == 2 && af.wire == 5) { av = bits_to_float(af.value); has_av = true; }
}
if (aname == "epsilon" && has_av) { eps = av; has_eps = true; }
}
}
if (op == "BatchNormalization" && has_eps && inputs.size() > 1)
bn_eps[inputs[1]] = eps;
out.op_type = std::move(op);
out.inputs = std::move(inputs);
}
} // namespace
std::string OnnxTensor::shapeString() const {
std::string s = "[";
for (size_t i = 0; i < shape.size(); ++i) {
if (i) s += ", ";
s += std::to_string(shape[i]);
}
return s + "]";
}
const OnnxTensor* OnnxFile::find(const std::string& name) const {
for (const OnnxTensor& t : initializers)
if (t.name == name) return &t;
return nullptr;
}
const OnnxNode* OnnxFile::producer(const std::string& tensor) const {
for (const OnnxNode& n : nodes)
for (const std::string& o : n.outputs)
if (o == tensor) return &n;
return nullptr;
}
std::unordered_map<std::string, std::string> OnnxFile::linearWeights() const {
std::unordered_set<std::string> inits;
for (const OnnxTensor& t : initializers) inits.insert(t.name);
std::unordered_map<std::string, std::string> out;
for (const OnnxNode& n : nodes) {
if (n.op_type != "Add" || n.inputs.size() != 2) continue;
for (int b = 0; b < 2; ++b) {
const std::string& bias = n.inputs[b];
if (!inits.count(bias)) continue;
const size_t dot = bias.rfind(".bias");
if (dot == std::string::npos || dot + 5 != bias.size()) continue;
const OnnxNode* mm = producer(n.inputs[1 - b]);
if (!mm || mm->op_type != "MatMul" || mm->inputs.size() != 2) continue;
// MatMul's second operand is the weight when it is an initializer.
if (inits.count(mm->inputs[1])) out[bias.substr(0, dot)] = mm->inputs[1];
}
}
return out;
}
OnnxFile read_onnx(const std::string& path) {
std::ifstream fin(path, std::ios::binary | std::ios::ate);
NN_CHECK((bool)fin, "cannot open '%s'", path.c_str());
const std::streamoff size = fin.tellg();
NN_CHECK(size > 16, "'%s' is %lld bytes; not an ONNX model", path.c_str(),
(long long)size);
fin.seekg(0);
std::vector<uint8_t> buf((size_t)size);
NN_CHECK((bool)fin.read(reinterpret_cast<char*>(buf.data()), size),
"cannot read '%s'", path.c_str());
Reader model{buf.data(), buf.data() + buf.size(), path.c_str()};
Reader graph{};
bool has_graph = false;
Field f;
while (next_field(model, f)) {
if (f.number == 7 && f.wire == 2) { graph = f.bytes; has_graph = true; }
}
NN_CHECK(has_graph, "'%s' has no graph; not an ONNX model", path.c_str());
OnnxFile out;
while (next_field(graph, f)) {
if (f.number == 5 && f.wire == 2) {
out.initializers.push_back(read_tensor_proto(f.bytes));
} else if (f.number == 1 && f.wire == 2) {
out.nodes.emplace_back();
scan_node(f.bytes, out.nodes.back(), out.bn_epsilon);
}
}
NN_CHECK(!out.initializers.empty(), "'%s' has no initializers", path.c_str());
size_t elems = 0;
for (const OnnxTensor& t : out.initializers) elems += t.data.size();
NN_LOG_DEBUG("[aliked] %s: %zu initializers, %.2f MB of weights\n", path.c_str(),
out.initializers.size(), (double)elems * 4.0 / 1e6);
return out;
}
} // namespace aliked
src/aliked/model/Onnx.h
+81
View File
@@ -0,0 +1,81 @@
#pragma once
// Just enough ONNX to read a checkpoint: the initializers, and the epsilon of
// every BatchNormalization node.
//
// We do not host or convert weights (docs/notes/aliked-port-plan.md). The
// artifact fetched is byte-for-byte the one COLMAP fetches, so a parity check
// against `colmap feature_extractor --FeatureExtraction.type ALIKED_N16ROT`
// compares two implementations of the same numbers rather than two
// checkpoints. The price is reading protobuf, and the price is small: the
// graph structure is hard-coded in AlikedModel.cpp, so nothing here has to
// understand a node, an operator, a type or a shape inference rule. It walks
// three nested messages by field number and skips everything else by wire
// type.
//
// ModelProto field 7 -> GraphProto
// GraphProto field 5 -> repeated TensorProto (initializers)
// field 1 -> repeated NodeProto (only for BN epsilon)
// TensorProto field 1 dims, 2 data_type, 8 name, 9 raw_data, 4 float_data
//
// This is deliberately NOT a general ONNX reader. It will not run a model, and
// it rejects rather than guesses: an initializer in a dtype we do not convert,
// or a truncated payload, throws with the tensor's name in the message.
#include <cstdint>
#include <string>
#include <unordered_map>
#include <vector>
namespace aliked {
// One initializer, converted to f32 on the host. ONNX shapes are PyTorch order
// already, which is also nn::Tensor's, so the shape is carried verbatim.
struct OnnxTensor {
std::string name;
std::vector<int64_t> shape;
std::vector<float> data;
int64_t numel() const {
int64_t n = 1;
for (int64_t d : shape) n *= d;
return n;
}
// "[16, 3, 3, 3]", for error messages.
std::string shapeString() const;
};
// Just enough of a node to resolve an anonymous initializer to the module it
// belongs to. LightGlue's export needs this: torch lowers nn.Linear to
// MatMul + Add, and only the Add's bias keeps its qualified name -- the weight
// arrives as "onnx::MatMul_2537". Walking Add -> MatMul recovers the pairing.
struct OnnxNode {
std::string op_type;
std::string name;
std::vector<std::string> inputs;
std::vector<std::string> outputs;
};
struct OnnxFile {
std::vector<OnnxTensor> initializers;
std::vector<OnnxNode> nodes;
// Keyed by the *scale* initializer's name -- BatchNormalization's input 1,
// which is the only name a caller folding BN into a conv already knows.
// Absent means the node did not spell epsilon out, i.e. the ONNX default.
std::unordered_map<std::string, float> bn_epsilon;
const OnnxTensor* find(const std::string& name) const;
// The node that produces `tensor`, or null.
const OnnxNode* producer(const std::string& tensor) const;
// For every `<module>.bias` initializer consumed by an Add whose other
// input comes from a MatMul, the name of that MatMul's weight. Keyed by
// the module prefix, so "transformers.0.self_attn.Wqkv" maps to the
// initializer holding its [in, out] matrix.
std::unordered_map<std::string, std::string> linearWeights() const;
};
// Throws nn::Error naming `path` on anything malformed.
OnnxFile read_onnx(const std::string& path);
} // namespace aliked
src/aliked/model/Weights.cpp
+300
View File
@@ -0,0 +1,300 @@
#include "aliked/model/Weights.h"
#include "aliked/Common.h"
#include "nn/vk/Memory.h"
#include "nn/vk/Stream.h"
#include <algorithm>
#include <cmath>
#include <cstdarg>
#include <cstdio>
#include <cstring>
namespace aliked {
namespace {
constexpr uint64_t kAlign = 256;
uint64_t align_up(uint64_t v, uint64_t a) { return (v + a - 1) / a * a; }
// A tensor staged on the host on its way to the device.
struct Staged {
std::string name;
std::vector<int64_t> shape;
std::vector<float> data;
};
const OnnxTensor& need(const OnnxFile& f, const std::string& name, const char* path) {
const OnnxTensor* t = f.find(name);
if (t) return *t;
// Name every miss with the file, because the usual cause is pointing
// --aliked-model at the wrong .onnx (the LightGlue one, say).
nn::fail("'%s' has no initializer '%s'; is this an ALIKED checkpoint?", path,
name.c_str());
}
void check_shape(const OnnxTensor& t, std::initializer_list<int64_t> want,
const char* path) {
bool ok = t.shape.size() == want.size();
if (ok) {
size_t i = 0;
for (int64_t w : want) {
if (w >= 0 && t.shape[i] != w) ok = false;
++i;
}
}
if (ok) return;
std::string ws = "[";
size_t i = 0;
for (int64_t w : want) {
if (i++) ws += ", ";
ws += (w < 0) ? std::string("*") : std::to_string(w);
}
ws += "]";
nn::fail("'%s': initializer '%s' is %s, expected %s", path, t.name.c_str(),
t.shapeString().c_str(), ws.c_str());
}
// Fold `conv` (shape [Co, Ci, kh, kw], no bias -- ALIKED's convs never have
// one) together with the BatchNorm that follows it into one conv plus a bias:
//
// s = gamma / sqrt(var + eps)
// w'[o] = w[o] * s[o]
// b'[o] = beta[o] - mean[o] * s[o]
//
// Appends both to `out`, named "<conv>.weight" and "<conv>.bias".
void fold_bn(const OnnxFile& f, const char* path, const std::string& conv_name,
const std::string& bn_prefix, std::vector<Staged>& out) {
const OnnxTensor& w = need(f, conv_name + ".weight", path);
const OnnxTensor& gamma = need(f, bn_prefix + ".weight", path);
const OnnxTensor& beta = need(f, bn_prefix + ".bias", path);
const OnnxTensor& mean = need(f, bn_prefix + ".running_mean", path);
const OnnxTensor& var = need(f, bn_prefix + ".running_var", path);
NN_CHECK(w.shape.size() == 4, "'%s': '%s' is %s, expected a 4-D conv kernel", path,
w.name.c_str(), w.shapeString().c_str());
const int64_t Co = w.shape[0];
const int64_t per_out = w.shape[1] * w.shape[2] * w.shape[3];
for (const OnnxTensor* t : {&gamma, &beta, &mean, &var})
NN_CHECK((int64_t)t->data.size() == Co,
"'%s': '%s' has %zu entries, expected %lld to match '%s'", path,
t->name.c_str(), t->data.size(), (long long)Co, w.name.c_str());
// The epsilon in the file, not PyTorch's default. ONNX's own default is
// 1e-5 and applies when the node did not spell the attribute out.
float eps = 1e-5f;
auto it = f.bn_epsilon.find(gamma.name);
if (it != f.bn_epsilon.end()) eps = it->second;
Staged sw{conv_name + ".weight", w.shape, w.data};
Staged sb{conv_name + ".bias", {Co}, std::vector<float>((size_t)Co, 0.0f)};
for (int64_t o = 0; o < Co; ++o) {
const float denom = std::sqrt(var.data[(size_t)o] + eps);
NN_CHECK(denom > 0.0f && std::isfinite(denom),
"'%s': '%s'[%lld] is %g; cannot fold BatchNorm", path, var.name.c_str(),
(long long)o, var.data[(size_t)o]);
const float s = gamma.data[(size_t)o] / denom;
float* row = sw.data.data() + (size_t)(o * per_out);
for (int64_t i = 0; i < per_out; ++i) row[i] *= s;
sb.data[(size_t)o] = beta.data[(size_t)o] - mean.data[(size_t)o] * s;
}
out.push_back(std::move(sw));
out.push_back(std::move(sb));
}
// Copy an initializer through unchanged.
void take(const OnnxFile& f, const char* path, const std::string& name,
std::vector<Staged>& out) {
const OnnxTensor& t = need(f, name, path);
out.push_back(Staged{name, t.shape, t.data});
}
void take_optional(const OnnxFile& f, const std::string& name, std::vector<Staged>& out) {
if (const OnnxTensor* t = f.find(name)) out.push_back(Staged{name, t->shape, t->data});
}
} // namespace
AlikedWeights::~AlikedWeights() {
if (blob_) vk::device_free(blob_);
}
void AlikedWeights::load(const std::string& onnx_path) {
NN_CHECK(!loaded_, "AlikedWeights::load called twice");
const char* path = onnx_path.c_str();
const OnnxFile file = read_onnx(onnx_path);
// ---- hyperparameters, from shapes ----
AlikedHparams hp;
{
const OnnxTensor& b1 = need(file, "block1.conv1.weight", path);
check_shape(b1, {-1, 3, 3, 3}, path);
hp.c1 = (int)b1.shape[0];
hp.c2 = (int)need(file, "block2.conv1.weight", path).shape[0];
hp.c3 = (int)need(file, "block3.conv1.regular_conv.weight", path).shape[0];
hp.c4 = (int)need(file, "block4.conv1.regular_conv.weight", path).shape[0];
const OnnxTensor& p1 = need(file, "conv1.weight", path);
check_shape(p1, {-1, -1, 1, 1}, path);
hp.dim4 = (int)p1.shape[0];
hp.dim = hp.dim4 * 4;
const OnnxTensor& agg = need(file, "desc_head.agg_weights", path);
NN_CHECK(agg.shape.size() == 3, "'%s': desc_head.agg_weights is %s, expected [M, C, C]",
path, agg.shapeString().c_str());
hp.M = (int)agg.shape[0];
hp.desc_dim = (int)agg.shape[2];
NN_CHECK(agg.shape[1] == hp.dim,
"'%s': desc_head.agg_weights is %s but the aggregated width is %d", path,
agg.shapeString().c_str(), hp.dim);
const OnnxTensor& oc = need(file, "desc_head.offset_conv.0.weight", path);
check_shape(oc, {2 * hp.M, hp.dim, -1, -1}, path);
NN_CHECK(oc.shape[2] == oc.shape[3], "'%s': desc_head.offset_conv.0 is not square",
path);
hp.K = (int)oc.shape[2];
// The deformable convs carry offsets only -- 2 per tap, no modulation
// mask. A 3*k*k export would need a different kernel, so refuse it
// rather than silently sampling with the wrong channel stride.
const OnnxTensor& off = need(file, "block3.conv1.offset_conv.weight", path);
NN_CHECK(off.shape.size() == 4 && off.shape[0] == 2 * off.shape[2] * off.shape[3],
"'%s': block3.conv1.offset_conv is %s; this port implements the "
"unmodulated (2*k*k offset channels) form only",
path, off.shapeString().c_str());
}
hp_ = hp;
// ---- stage every tensor the forward pass asks for ----
std::vector<Staged> staged;
staged.reserve(48);
// block1 is a ConvBlock, block2..4 are ResBlocks; blocks 3 and 4 use
// deformable convs, whose *offset* conv has its own bias and no BN.
fold_bn(file, path, "block1.conv1", "block1.bn1", staged);
fold_bn(file, path, "block1.conv2", "block1.bn2", staged);
fold_bn(file, path, "block2.conv1", "block2.bn1", staged);
fold_bn(file, path, "block2.conv2", "block2.bn2", staged);
take(file, path, "block2.downsample.weight", staged);
take(file, path, "block2.downsample.bias", staged);
for (int b = 3; b <= 4; ++b) {
char buf[64];
for (int c = 1; c <= 2; ++c) {
std::snprintf(buf, sizeof buf, "block%d.conv%d.regular_conv", b, c);
char bn[64];
std::snprintf(bn, sizeof bn, "block%d.bn%d", b, c);
fold_bn(file, path, buf, bn, staged);
std::snprintf(buf, sizeof buf, "block%d.conv%d.offset_conv.weight", b, c);
take(file, path, buf, staged);
std::snprintf(buf, sizeof buf, "block%d.conv%d.offset_conv.bias", b, c);
take(file, path, buf, staged);
}
std::snprintf(buf, sizeof buf, "block%d.downsample.weight", b);
take(file, path, buf, staged);
std::snprintf(buf, sizeof buf, "block%d.downsample.bias", b);
take(file, path, buf, staged);
}
for (int i = 1; i <= 4; ++i) {
char buf[32];
std::snprintf(buf, sizeof buf, "conv%d.weight", i);
take(file, path, buf, staged);
}
for (int i : {0, 2, 4, 6}) {
char buf[32];
std::snprintf(buf, sizeof buf, "score_head.%d.weight", i);
take(file, path, buf, staged);
// The score head's convs have no bias in the released checkpoints;
// read one if some future export grows it rather than assuming.
std::snprintf(buf, sizeof buf, "score_head.%d.bias", i);
take_optional(file, buf, staged);
}
take(file, path, "desc_head.offset_conv.0.weight", staged);
take(file, path, "desc_head.offset_conv.0.bias", staged);
take(file, path, "desc_head.offset_conv.2.weight", staged);
take(file, path, "desc_head.offset_conv.2.bias", staged);
take(file, path, "desc_head.sf_conv.weight", staged);
take_optional(file, "desc_head.sf_conv.bias", staged);
// The descriptor aggregation is einsum('ncp,pcd->nd'): for each of the M
// sample positions, a [C, D] matrix applied to that position's features
// and summed. That is M ordinary matmuls -- which is how the forward pass
// runs it, on the tuned GEMM -- except that `linear` wants its weight as
// [out_features, in_features] and the checkpoint stores [in, out]. So the
// transpose happens once, here, rather than M times per image in a kernel.
{
const OnnxTensor& agg = need(file, "desc_head.agg_weights", path);
const int64_t M = agg.shape[0], C = agg.shape[1], D = agg.shape[2];
Staged t{"desc_head.agg_weights_t", {M, D, C}, std::vector<float>(agg.data.size())};
for (int64_t p = 0; p < M; ++p)
for (int64_t c = 0; c < C; ++c)
for (int64_t d = 0; d < D; ++d)
t.data[(size_t)((p * D + d) * C + c)] =
agg.data[(size_t)((p * C + c) * D + d)];
staged.push_back(std::move(t));
}
// ---- one allocation, 256-byte sub-alignment ----
uint64_t total = 0;
for (const Staged& s : staged) total = align_up(total, kAlign) + s.data.size() * 4;
NN_CHECK(total > 0, "'%s': nothing to upload", path);
blob_ = vk::device_alloc(total, "aliked-weights");
device_bytes_ = total;
uint64_t off = 0;
for (const Staged& s : staged) {
off = align_up(off, kAlign);
const nn::DevicePtr ptr = blob_ + off;
vk::Stream::get().upload(ptr, s.data.data(), s.data.size() * 4);
off += s.data.size() * 4;
nn::Tensor t;
t.ptr = ptr;
t.dtype = nn::DType::F32;
t.ndim = (int32_t)std::min<size_t>(s.shape.size(), 4);
NN_CHECK(s.shape.size() <= 4, "'%s': '%s' has rank %zu; nn::Tensor holds 4", path,
s.name.c_str(), s.shape.size());
for (int i = 0; i < t.ndim; ++i) t.shape[i] = s.shape[(size_t)i];
if (t.ndim == 0) { t.ndim = 1; t.shape[0] = 1; }
tensors_[s.name] = t;
}
vk::Stream::get().sync();
path_ = onnx_path;
loaded_ = true;
NN_LOG_INFO("[aliked] %s: c=(%d,%d,%d,%d) dim=%d K=%d M=%d desc=%d, %zu tensors, "
"%.2f MB on device\n",
onnx_path.c_str(), hp.c1, hp.c2, hp.c3, hp.c4, hp.dim, hp.K, hp.M,
hp.desc_dim, tensors_.size(), (double)total / 1e6);
}
nn::Tensor AlikedWeights::get(const std::string& name) const {
auto it = tensors_.find(name);
if (it != tensors_.end()) return it->second;
// Near misses: the failure is almost always a name that drifted, and
// printing the neighbourhood turns a five-minute hunt into a glance.
std::vector<std::string> near;
for (const auto& kv : tensors_) {
const std::string& k = kv.first;
const size_t dot = name.find('.');
if (dot != std::string::npos && k.compare(0, dot, name, 0, dot) == 0)
near.push_back(k);
}
std::sort(near.begin(), near.end());
std::string hint;
for (size_t i = 0; i < near.size() && i < 8; ++i) hint += "\n " + near[i];
nn::fail("no weight named '%s'%s%s", name.c_str(),
hint.empty() ? "" : " (did you mean one of these?)", hint.c_str());
}
nn::Tensor AlikedWeights::getf(const char* fmt, ...) const {
char buf[256];
va_list ap;
va_start(ap, fmt);
std::vsnprintf(buf, sizeof buf, fmt, ap);
va_end(ap);
return get(buf);
}
} // namespace aliked
src/aliked/model/Weights.h
+73
View File
@@ -0,0 +1,73 @@
#pragma once
// The ALIKED checkpoint: an .onnx file in, named device tensors out.
//
// Two things happen between those, and both are here rather than in the
// forward pass because both are properties of the *file*:
//
// * BatchNorm folding. The export did not fold it (8 BatchNormalization
// nodes, running_mean / running_var present as initializers), so every
// conv-then-BN pair becomes one conv with a bias. The epsilon comes from
// the node attribute, not from PyTorch's default -- Onnx.cpp reads it.
// * Hyperparameters. `aliked-n16rot` and `aliked-n32` have byte-identical
// graph structure and differ only in M, the number of SDDH sample
// positions. So M is read from desc_head.agg_weights, not from the id, and
// one code path serves both. Same for every channel count.
//
// Everything stays f32 on the device. The whole checkpoint is 2.7 MB; what
// costs memory in this model is the full-resolution activations, and halving
// 2.7 MB against that is not worth a second numeric path (docs/notes/
// aliked-port-plan.md).
#include "aliked/model/Onnx.h"
#include "nn/Tensor.h"
#include <string>
#include <unordered_map>
#include <vector>
namespace aliked {
// Read off the checkpoint's tensor shapes; see the class comment.
struct AlikedHparams {
int c1 = 0, c2 = 0, c3 = 0, c4 = 0; // encoder block widths
int dim = 0; // aggregated feature width (4 * dim4)
int dim4 = 0; // per-block projection width, dim / 4
int K = 0; // SDDH patch size (3)
int M = 0; // SDDH sample positions (16 or 32)
int desc_dim = 0; // descriptor width (128)
};
class AlikedWeights {
public:
AlikedWeights() = default;
~AlikedWeights();
AlikedWeights(const AlikedWeights&) = delete;
AlikedWeights& operator=(const AlikedWeights&) = delete;
// Parses, validates, folds and uploads. Throws nn::Error naming the tensor
// on any missing or unexpectedly shaped weight -- a checkpoint that is not
// ALIKED must fail here with a sentence, not later with a fault.
void load(const std::string& onnx_path);
bool loaded() const { return loaded_; }
const AlikedHparams& hparams() const { return hp_; }
const std::string& path() const { return path_; }
uint64_t deviceBytes() const { return device_bytes_; }
// Throws when absent, listing near misses: a typo in a weight name is
// otherwise a null tensor that faults somewhere else entirely.
nn::Tensor get(const std::string& name) const;
// printf-style, for the per-block names ("block%d.conv1.weight").
nn::Tensor getf(const char* fmt, ...) const;
bool has(const std::string& name) const { return tensors_.count(name) != 0; }
private:
std::unordered_map<std::string, nn::Tensor> tensors_;
AlikedHparams hp_;
std::string path_;
nn::DevicePtr blob_ = 0;
uint64_t device_bytes_ = 0;
bool loaded_ = false;
};
} // namespace aliked
src/aliked/shaders/aliked.slang
+426
View File
@@ -0,0 +1,426 @@
// ALIKED's model-specific kernels: the detector's NMS and sub-pixel
// refinement, and the two coordinate conversions its descriptor head needs.
//
// Everything else in the forward pass is a general op from nn/ -- convolution,
// deformable convolution, pooling, resize, grid sample, GEMM, L2 normalize.
// What is left here is what could not be anything else: a suppression rule, a
// soft-argmax, and the layout conversions between them.
#include "_common.slang"
// Spec ids 0..1.
[SpecializationConstant] const int kInputF16 = 0;
[SpecializationConstant] const int kAct = 0;
static const int WG = 256;
// ---------------------------------------------------------------------------
// Non-maximum suppression
//
// A faithful port of ALIKED's `simple_nms`, which is NOT "score == maxpool":
//
// max_mask = scores == max_pool(scores)
// twice:
// supp_mask = max_pool(max_mask) > 0
// supp_scores = where(supp_mask, 0, scores)
// max_mask |= (supp_scores == max_pool(supp_scores)) & ~supp_mask
//
// The two extra rounds recover a peak that a *neighbouring* peak masked, which
// on a smooth score map is a few percent of all keypoints -- not a rounding
// difference, a different feature set. The five max_pool calls this expands to
// are exactly the five MaxPool nodes in the exported graph, which is how the
// port was checked against it.
//
// Split across two entry points per round because supp_mask has to be
// materialized: computing whether a *neighbour* is suppressed needs the whole
// window of max_mask, so it cannot be folded into the comparison that reads it.
// ---------------------------------------------------------------------------
struct NmsParams {
uint* mask; // [H, W] u32, 0/1 -- read and written across rounds
uint* supp; // [H, W] u32, 0/1 scratch
uint* scores; // [H, W] f32
uint H;
uint W;
int radius;
uint groups_per_row;
};
// max_mask = (scores == max over the (2r+1)^2 window)
[shader("compute")]
[numthreads(WG, 1, 1)]
void nms_init(uint3 gid : SV_GroupID, uint3 tid : SV_GroupThreadID, uniform NmsParams p) {
uint i = fold_index(gid, tid, p.groups_per_row, uint(WG));
if (i >= p.H * p.W) return;
int x = int(i % p.W), y = int(i / p.W);
float s = asfloat(p.scores[i]);
float best = -3.402823466e+38f;
for (int dy = -p.radius; dy <= p.radius; ++dy) {
int yy = y + dy;
if (yy < 0 || yy >= int(p.H)) continue;
for (int dx = -p.radius; dx <= p.radius; ++dx) {
int xx = x + dx;
if (xx < 0 || xx >= int(p.W)) continue;
best = max(best, asfloat(p.scores[uint(yy) * p.W + uint(xx)]));
}
}
p.mask[i] = (s == best) ? 1u : 0u;
}
// supp = max_pool(max_mask) > 0, i.e. "some accepted peak lies within radius".
[shader("compute")]
[numthreads(WG, 1, 1)]
void nms_suppress(uint3 gid : SV_GroupID, uint3 tid : SV_GroupThreadID,
uniform NmsParams p) {
uint i = fold_index(gid, tid, p.groups_per_row, uint(WG));
if (i >= p.H * p.W) return;
int x = int(i % p.W), y = int(i / p.W);
uint any = 0u;
for (int dy = -p.radius; dy <= p.radius && any == 0u; ++dy) {
int yy = y + dy;
if (yy < 0 || yy >= int(p.H)) continue;
for (int dx = -p.radius; dx <= p.radius; ++dx) {
int xx = x + dx;
if (xx < 0 || xx >= int(p.W)) continue;
if (p.mask[uint(yy) * p.W + uint(xx)] != 0u) { any = 1u; break; }
}
}
p.supp[i] = any;
}
// max_mask |= (supp_scores == max_pool(supp_scores)) & ~supp
[shader("compute")]
[numthreads(WG, 1, 1)]
void nms_recover(uint3 gid : SV_GroupID, uint3 tid : SV_GroupThreadID,
uniform NmsParams p) {
uint i = fold_index(gid, tid, p.groups_per_row, uint(WG));
if (i >= p.H * p.W) return;
if (p.supp[i] != 0u) return; // ~supp_mask gates the whole term
int x = int(i % p.W), y = int(i / p.W);
// supp_scores is `where(supp, 0, scores)`, read rather than materialized.
float s = asfloat(p.scores[i]);
float best = -3.402823466e+38f;
for (int dy = -p.radius; dy <= p.radius; ++dy) {
int yy = y + dy;
if (yy < 0 || yy >= int(p.H)) continue;
for (int dx = -p.radius; dx <= p.radius; ++dx) {
int xx = x + dx;
if (xx < 0 || xx >= int(p.W)) continue;
uint j = uint(yy) * p.W + uint(xx);
best = max(best, (p.supp[j] != 0u) ? 0.0f : asfloat(p.scores[j]));
}
}
if (s == best) p.mask[i] = 1u;
}
// ---------------------------------------------------------------------------
// Candidate compaction
// ---------------------------------------------------------------------------
struct CollectParams {
uint* out; // [cap, 3]: x, y (u32) and the score's bits
uint* counter; // [1]
uint* mask;
uint* scores;
uint H;
uint W;
int border; // DKD zeroes this many rows/cols at the top and left,
// and border-1 at the bottom and right (see below)
float min_score; // peaks below this never become candidates
uint cap;
uint groups_per_row;
};
// Appends every surviving peak, with its integer position and score.
//
// The border rule is ALIKED's and is asymmetric -- it zeroes rows [0, r] and
// columns [0, r] but only rows [h-r, h) and columns [w-r, w). So a keypoint
// satisfies r < y < h-r, which is also exactly the condition under which its
// (2r+1)^2 refinement patch is fully inside the map. The refinement kernel
// relies on that and does no bounds testing.
[shader("compute")]
[numthreads(WG, 1, 1)]
void nms_collect(uint3 gid : SV_GroupID, uint3 tid : SV_GroupThreadID,
uniform CollectParams p) {
uint i = fold_index(gid, tid, p.groups_per_row, uint(WG));
if (i >= p.H * p.W) return;
if (p.mask[i] == 0u) return;
// The threshold applies to the PEAK's own score here, and again to the
// sub-pixel-refined score after refinement. Both, because that is what the
// reference does: a peak scoring just under the threshold whose refined
// score lands just over it is dropped, not kept. Filtering here also keeps
// the candidate list bounded on a textured image.
if (asfloat(p.scores[i]) < p.min_score) return;
int x = int(i % p.W), y = int(i / p.W);
if (x <= p.border || y <= p.border) return;
if (x >= int(p.W) - p.border || y >= int(p.H) - p.border) return;
uint slot;
InterlockedAdd(p.counter[0], 1u, slot);
if (slot >= p.cap) return; // the host reads the counter and warns
p.out[slot * 3u + 0u] = uint(x);
p.out[slot * 3u + 1u] = uint(y);
p.out[slot * 3u + 2u] = p.scores[i];
}
// ---------------------------------------------------------------------------
// Sub-pixel refinement
// ---------------------------------------------------------------------------
struct RefineParams {
float* out_xy; // [N, 2] refined pixel coordinates, (x, y)
float* out_score; // [N]
uint* centers; // [N, 2] i32 integer peak positions, (x, y)
uint* scores; // [H, W] f32
uint H;
uint W;
uint N;
int radius;
float temperature;
uint groups_per_row;
};
// ALIKED's DKD soft-argmax: over the (2r+1)^2 patch of the *score map* around
// the integer peak,
//
// e = exp((patch - max(patch)) / T)
// d = sum(e * grid) / sum(e) grid = the (dx, dy) offsets, -r..r
// xy = peak + d
//
// and the reported score is the score map resampled *bilinearly at the refined
// position*, not the peak's own value -- COLMAP then thresholds on that one, so
// using the peak value here would filter a different set.
//
// One thread per keypoint: the patch is 25 values and there are a few thousand
// of them, so a workgroup-per-keypoint reduction would leave most lanes idle.
[shader("compute")]
[numthreads(WG, 1, 1)]
void refine_keypoints(uint3 gid : SV_GroupID, uint3 tid : SV_GroupThreadID,
uniform RefineParams p) {
uint n = fold_index(gid, tid, p.groups_per_row, uint(WG));
if (n >= p.N) return;
int cx = int(p.centers[2u * n]);
int cy = int(p.centers[2u * n + 1u]);
// The border rule guarantees the patch is inside; no clamping, so a
// violation would be a bug worth crashing on rather than hiding.
float mx = -3.402823466e+38f;
for (int dy = -p.radius; dy <= p.radius; ++dy)
for (int dx = -p.radius; dx <= p.radius; ++dx)
mx = max(mx, asfloat(p.scores[uint(cy + dy) * p.W + uint(cx + dx)]));
float sum = 0.0f, sx = 0.0f, sy = 0.0f;
for (int dy = -p.radius; dy <= p.radius; ++dy)
for (int dx = -p.radius; dx <= p.radius; ++dx) {
float e = exp((asfloat(p.scores[uint(cy + dy) * p.W + uint(cx + dx)]) - mx) /
p.temperature);
sum += e;
sx += e * float(dx);
sy += e * float(dy);
}
float inv = 1.0f / max(sum, 1e-20f);
float rx = float(cx) + sx * inv;
float ry = float(cy) + sy * inv;
p.out_xy[2u * n] = rx;
p.out_xy[2u * n + 1u] = ry;
// grid_sample(align_corners=True) at the refined point: in pixel space
// that is a plain bilinear fetch, and the point is interior by
// construction, so clamped and zero-padded sampling agree here.
p.out_score[n] = sample_bilinear(p.scores, false, int(p.H), int(p.W), 1, 0, ry, rx);
}
// ---------------------------------------------------------------------------
// Descriptor-head coordinate conversions
// ---------------------------------------------------------------------------
struct CenterParams {
uint* out; // [N, 2] i32
uint* xy; // [N, 2] f32 pixel coordinates
uint N;
uint groups_per_row;
};
// The patch centre SDDH gathers around is `kptsi_wh.long()` -- a truncation of
// the *refined* sub-pixel position, which is not always the integer peak the
// refinement started from. Doing this on the host would mean a round trip
// between the refinement and the descriptor head for no reason.
[shader("compute")]
[numthreads(WG, 1, 1)]
void to_int_centers(uint3 gid : SV_GroupID, uint3 tid : SV_GroupThreadID,
uniform CenterParams p) {
uint i = fold_index(gid, tid, p.groups_per_row, uint(WG));
if (i >= p.N * 2u) return;
// Coordinates are non-negative here (the border rule keeps them >= 1), so
// torch's truncation and a floor agree.
p.out[i] = uint(int(asfloat(p.xy[i])));
}
struct SddhPosParams {
float* out; // [N * M, 2] normalized (gx, gy) for grid_sample
uint* xy; // [N, 2] f32 keypoint pixel coordinates
uint* offsets; // [N, 2 * M] f32, raw offset-head output
uint N;
uint M;
uint W;
uint H;
float max_offset;
uint groups_per_row;
};
// SDDH's sampling positions.
//
// The offset head emits 2*M channels which ALIKED reads as `view(N, 2, M)`:
// the FIRST M channels are the x components and the SECOND M the y ones. That
// is not the layout of the deformable convolution's offsets in the same
// network, which are (dy, dx) interleaved per tap. Two conventions, one model;
// swapping either produces descriptors that look normal and match badly.
//
// Positions are then normalized on the align_corners=True convention --
// 2*pos/(w-1, h-1) - 1 -- because that is what ALIKED's grid_sample call uses.
[shader("compute")]
[numthreads(WG, 1, 1)]
void sddh_positions(uint3 gid : SV_GroupID, uint3 tid : SV_GroupThreadID,
uniform SddhPosParams p) {
uint i = fold_index(gid, tid, p.groups_per_row, uint(WG));
if (i >= p.N * p.M) return;
uint m = i % p.M;
uint n = i / p.M;
float ox = asfloat(p.offsets[n * 2u * p.M + m]);
float oy = asfloat(p.offsets[n * 2u * p.M + p.M + m]);
ox = clamp(ox, -p.max_offset, p.max_offset);
oy = clamp(oy, -p.max_offset, p.max_offset);
float px = asfloat(p.xy[2u * n]) + ox;
float py = asfloat(p.xy[2u * n + 1u]) + oy;
p.out[2u * i] = 2.0f * px / (float(p.W) - 1.0f) - 1.0f;
p.out[2u * i + 1u] = 2.0f * py / (float(p.H) - 1.0f) - 1.0f;
}
// ---------------------------------------------------------------------------
// LightGlue's assignment tail
//
// The last transformer layer produces a similarity matrix `sim[N0, N1]` and a
// matchability logit per keypoint. LightGlue turns those into an assignment
// with `sigmoid_log_double_softmax`:
//
// score[i][j] = logsoftmax_j(sim)[i][j] + logsoftmax_i(sim)[i][j]
// + logsigmoid(z0[i]) + logsigmoid(z1[j])
//
// and then keeps a pair only if it is the argmax of BOTH its row and its
// column, with exp(score) over a threshold.
//
// Materializing `score` would be a second N0 x N1 matrix -- 16 MB at 2048
// keypoints each, written and read once. It is not needed: every term is
// separable into a row constant and a column constant, so the two log-sum-exps
// are vectors and the arg-max reads `sim` with those applied on the fly. What
// is left is three passes over the matrix instead of five, and no second
// matrix.
// ---------------------------------------------------------------------------
struct AssignParams {
float* out_val; // [R] the row's best score, before the row-constant terms
uint* out_idx; // [R] its column
uint* sim; // [N0, N1] f32
uint* bias; // [C] added to every element of a row before the max
uint R; // rows of the logical matrix this pass reduces over
uint C; // columns
uint stride; // element stride between logical rows of `sim`
uint step; // element step between logical columns
// What `sim` is multiplied by before the bias. TWO for the arg-max, not
// one: the score is logsoftmax_rows + logsoftmax_cols, and each of those
// carries its own copy of sim -- so the element being maximized is
// 2*sim[i][j] - lse_row[i] - lse_col[j] + ls0[i] + ls1[j]. Dropping the
// factor changes which column wins, not just the value, and leaves every
// score far enough below zero that exp() is 0 and nothing matches.
float scale;
uint groups_per_row;
};
// `stride`/`step` are what make one kernel serve both directions: reducing
// over columns of sim is (stride = N1, step = 1), over rows is
// (stride = 1, step = N1). Reading a column is a strided gather either way --
// the matrix is too big to transpose for the sake of one pass over it.
groupshared float g_val[WG];
groupshared uint g_idx[WG];
// logsumexp over each logical row, in one workgroup per row.
[shader("compute")]
[numthreads(WG, 1, 1)]
void assign_logsumexp(uint3 gid : SV_GroupID, uint3 tid : SV_GroupThreadID,
uniform AssignParams p) {
uint row = gid.y * p.groups_per_row + gid.x;
uint t = tid.x;
// Two-pass max-then-sum, both as groupshared trees so wave32 and wave64
// devices agree bit for bit.
float m = -3.402823466e+38f;
if (row < p.R)
for (uint c = t; c < p.C; c += uint(WG))
m = max(m, asfloat(p.sim[row * p.stride + c * p.step]));
g_val[t] = m;
GroupMemoryBarrierWithGroupSync();
for (uint s = uint(WG) / 2u; s > 0u; s >>= 1) {
if (t < s) g_val[t] = max(g_val[t], g_val[t + s]);
GroupMemoryBarrierWithGroupSync();
}
float mx = g_val[0];
GroupMemoryBarrierWithGroupSync();
float sum = 0.0f;
if (row < p.R)
for (uint c = t; c < p.C; c += uint(WG))
sum += exp(asfloat(p.sim[row * p.stride + c * p.step]) - mx);
g_val[t] = sum;
GroupMemoryBarrierWithGroupSync();
for (uint s = uint(WG) / 2u; s > 0u; s >>= 1) {
if (t < s) g_val[t] += g_val[t + s];
GroupMemoryBarrierWithGroupSync();
}
if (row < p.R && t == 0u) p.out_val[row] = mx + log(g_val[0]);
}
// argmax over each logical row of (sim + bias), where `bias` carries every
// term that depends on the column alone.
[shader("compute")]
[numthreads(WG, 1, 1)]
void assign_argmax(uint3 gid : SV_GroupID, uint3 tid : SV_GroupThreadID,
uniform AssignParams p) {
uint row = gid.y * p.groups_per_row + gid.x;
uint t = tid.x;
float best = -3.402823466e+38f;
uint bi = 0u;
if (row < p.R)
for (uint c = t; c < p.C; c += uint(WG)) {
float v = p.scale * asfloat(p.sim[row * p.stride + c * p.step]) +
asfloat(p.bias[c]);
// Strictly greater, so the lowest index wins a tie -- the reduction
// order below is not the column order, and a mutual-nearest test
// that disagreed with itself run to run would make the whole
// reconstruction irreproducible.
if (v > best) { best = v; bi = c; }
}
g_val[t] = best;
g_idx[t] = bi;
GroupMemoryBarrierWithGroupSync();
for (uint s = uint(WG) / 2u; s > 0u; s >>= 1) {
if (t < s) {
const bool take = g_val[t + s] > g_val[t] ||
(g_val[t + s] == g_val[t] && g_idx[t + s] < g_idx[t]);
if (take) { g_val[t] = g_val[t + s]; g_idx[t] = g_idx[t + s]; }
}
GroupMemoryBarrierWithGroupSync();
}
if (row < p.R && t == 0u) {
p.out_val[row] = g_val[0];
p.out_idx[row] = g_idx[0];
}
}
src/aliked/tests/aliked_test.cpp
+444
View File
@@ -0,0 +1,444 @@
// ALIKED: the checkpoint reader, and (once it lands) the forward pass.
//
// The reader gate is deliberately strict about *shapes*, not about values: we
// do not own these weights and cannot embed a golden copy of them, so what
// this can check is that the file parses, that every tensor the forward pass
// will ask for exists at the width the rest of the model assumes, and that the
// two released variants agree everywhere except M.
//
// Needs a checkpoint. With none on disk and no --fetch it SKIPS rather than
// fails: a build machine without network access must still be able to run the
// test suite.
//
// aliked_test # use whatever is already cached
// aliked_test --fetch # download from COLMAP's releases if missing
// aliked_test path/to.onnx # a specific file
#include "aliked/Aliked.h"
#include "aliked/Common.h"
#include "aliked/model/Fetch.h"
#include "aliked/model/LightGlue.h"
#include "aliked/model/Weights.h"
#include "nn/core/Log.h"
#include "nn/io/Image.h"
#include "nn/vk/Context.h"
#include "nn/vk/Memory.h"
#include "nn/vk/Pipelines.h"
#include "nn/vk/Stream.h"
#include <algorithm>
#include <cmath>
#include <cstdarg>
#include <cstdlib>
#include <cstring>
#include <cstdio>
#include <filesystem>
#include <string>
#include <vector>
namespace fs = std::filesystem;
using namespace aliked;
namespace {
int g_failures = 0;
int g_checks = 0;
void check(bool ok, const char* fmt, ...) {
++g_checks;
if (ok) return;
++g_failures;
char buf[512];
va_list ap;
va_start(ap, fmt);
std::vsnprintf(buf, sizeof buf, fmt, ap);
va_end(ap);
std::printf(" FAIL %s\n", buf);
}
void expect_shape(const AlikedWeights& w, const char* name,
std::vector<int64_t> want) {
nn::Tensor t;
try {
t = w.get(name);
} catch (const std::exception& e) {
check(false, "%s: %s", name, e.what());
return;
}
bool ok = t.ndim == (int32_t)want.size();
for (size_t i = 0; ok && i < want.size(); ++i)
if (want[i] >= 0 && t.shape[i] != want[i]) ok = false;
if (!ok) {
std::string got = "[";
for (int i = 0; i < t.ndim; ++i)
got += (i ? ", " : "") + std::to_string(t.shape[i]);
got += "]";
std::string ws = "[";
for (size_t i = 0; i < want.size(); ++i)
ws += (i ? ", " : "") + (want[i] < 0 ? std::string("*") : std::to_string(want[i]));
ws += "]";
check(false, "%s is %s, expected %s", name, got.c_str(), ws.c_str());
} else {
++g_checks;
}
}
// Everything AlikedModel's forward pass will look up, at the widths it assumes.
void check_checkpoint(const std::string& path) {
std::printf("\n%s\n", path.c_str());
// Whenever the file we were handed is recognisably one of the released
// artifacts, hash it. This is the only place SHA-256 gets exercised on
// real data -- everywhere else it runs on the download path, where a bug
// would show up as "the file is always corrupt" long after the fact.
for (const char* id : {"aliked-n16rot", "aliked-n32", "aliked-lightglue"}) {
const ModelSource* src = find_model_source(id);
if (!src || path.find(src->file) == std::string::npos) continue;
const std::string got = sha256_file(path);
check(got == src->sha256, "%s hashes to %s, expected %s", src->file, got.c_str(),
src->sha256);
break;
}
AlikedWeights w;
w.load(path);
const AlikedHparams& h = w.hparams();
std::printf(" c=(%d,%d,%d,%d) dim=%d dim4=%d K=%d M=%d desc=%d %.2f MB\n", h.c1,
h.c2, h.c3, h.c4, h.dim, h.dim4, h.K, h.M, h.desc_dim,
(double)w.deviceBytes() / 1e6);
// The published variants: n16 / n16rot / n32 all share these.
check(h.c1 == 16 && h.c2 == 32 && h.c3 == 64 && h.c4 == 128,
"unexpected block widths (%d,%d,%d,%d)", h.c1, h.c2, h.c3, h.c4);
check(h.dim == 128 && h.dim4 == 32, "unexpected aggregated width %d", h.dim);
check(h.desc_dim == 128, "unexpected descriptor width %d", h.desc_dim);
check(h.K == 3, "unexpected SDDH patch %d", h.K);
check(h.M == 16 || h.M == 32, "unexpected SDDH position count %d", h.M);
// Encoder. block1 is a ConvBlock (no shortcut projection); 2..4 are
// ResBlocks; 3 and 4 are deformable.
expect_shape(w, "block1.conv1.weight", {h.c1, 3, 3, 3});
expect_shape(w, "block1.conv1.bias", {h.c1});
expect_shape(w, "block1.conv2.weight", {h.c1, h.c1, 3, 3});
expect_shape(w, "block1.conv2.bias", {h.c1});
check(!w.has("block1.downsample.weight"),
"block1 has a downsample; it should be a ConvBlock");
const int cin[5] = {0, 0, 16, 32, 64};
const int cout[5] = {0, 0, 32, 64, 128};
for (int b = 2; b <= 4; ++b) {
char n[64];
const bool deform = b >= 3;
const char* suffix = deform ? ".regular_conv" : "";
std::snprintf(n, sizeof n, "block%d.conv1%s.weight", b, suffix);
expect_shape(w, n, {cout[b], cin[b], 3, 3});
std::snprintf(n, sizeof n, "block%d.conv1%s.bias", b, suffix);
expect_shape(w, n, {cout[b]});
std::snprintf(n, sizeof n, "block%d.conv2%s.weight", b, suffix);
expect_shape(w, n, {cout[b], cout[b], 3, 3});
std::snprintf(n, sizeof n, "block%d.downsample.weight", b);
expect_shape(w, n, {cout[b], cin[b], 1, 1});
std::snprintf(n, sizeof n, "block%d.downsample.bias", b);
expect_shape(w, n, {cout[b]});
if (deform) {
// 2 offsets per tap, no modulation mask -- Weights.cpp refuses the
// 3*k*k form, and this is the assertion that says why.
std::snprintf(n, sizeof n, "block%d.conv1.offset_conv.weight", b);
expect_shape(w, n, {18, cin[b], 3, 3});
std::snprintf(n, sizeof n, "block%d.conv2.offset_conv.weight", b);
expect_shape(w, n, {18, cout[b], 3, 3});
}
}
// Aggregation and heads.
const int ci[5] = {0, 16, 32, 64, 128};
for (int i = 1; i <= 4; ++i) {
char n[32];
std::snprintf(n, sizeof n, "conv%d.weight", i);
expect_shape(w, n, {h.dim4, ci[i], 1, 1});
}
expect_shape(w, "score_head.0.weight", {8, h.dim, 1, 1});
expect_shape(w, "score_head.2.weight", {4, 8, 3, 3});
expect_shape(w, "score_head.4.weight", {4, 4, 3, 3});
expect_shape(w, "score_head.6.weight", {1, 4, 3, 3});
expect_shape(w, "desc_head.offset_conv.0.weight", {2 * h.M, h.dim, h.K, h.K});
expect_shape(w, "desc_head.offset_conv.0.bias", {2 * h.M});
expect_shape(w, "desc_head.offset_conv.2.weight", {2 * h.M, 2 * h.M, 1, 1});
expect_shape(w, "desc_head.offset_conv.2.bias", {2 * h.M});
expect_shape(w, "desc_head.sf_conv.weight", {h.desc_dim, h.dim, 1, 1});
// Stored transposed at load: the aggregation runs as M matmuls on the
// tuned GEMM, which wants [out_features, in_features].
expect_shape(w, "desc_head.agg_weights_t", {h.M, h.desc_dim, h.dim});
// A name that is not there must say so, and say what is nearby.
bool threw = false;
try {
(void)w.get("block1.conv3.weight");
} catch (const std::exception&) {
threw = true;
}
check(threw, "a missing weight name did not throw");
}
// Extract from one image and write the result where a comparison script can
// read it. The format is deliberately trivial -- a header and two float
// blocks -- because its only consumer is tools/aliked/compare_colmap.py, which
// reads COLMAP's SQLite database on the other side.
//
// char[8] "ALIKEDFT" u32 version=1 i32 width, height u32 count, dim
// count * { f32 x, y, score }
// count * dim f32 descriptors
void run_extraction(const std::string& model, const std::string& image_path,
const std::string& out_path, const ExtractOptions& opts,
int max_image_size) {
nn::Image img = nn::load_image(image_path);
check(!img.empty(), "cannot read image %s", image_path.c_str());
if (img.empty()) return;
// COLMAP downscales to EffMaxImageSize() before extracting -- 1600 for
// ALIKED. Comparing against its output means doing the same, and doing it
// the same way: a box-filtered halving is not what it does, but at the
// integer ratios these tests use the difference is not what is under test.
if (max_image_size > 0 &&
std::max(img.width, img.height) > max_image_size) {
const double s = (double)max_image_size / std::max(img.width, img.height);
const int nw = std::max(1, (int)(img.width * s));
const int nh = std::max(1, (int)(img.height * s));
std::vector<uint8_t> dst((size_t)nw * nh * 3);
for (int y = 0; y < nh; ++y)
for (int x = 0; x < nw; ++x) {
const int sy = std::min(img.height - 1, (int)((y + 0.5) / s));
const int sx = std::min(img.width - 1, (int)((x + 0.5) / s));
std::memcpy(&dst[((size_t)y * nw + x) * 3],
&img.data[((size_t)sy * img.width + sx) * 3], 3);
}
img.data.swap(dst);
img.width = nw;
img.height = nh;
}
Extractor ex;
ex.load(model);
const double t0 = nn::now_ms();
Features f = ex.extract(img.data.data(), img.width, img.height, opts);
const double ms = nn::now_ms() - t0;
double min_s = 1e9, max_s = -1e9, mean_norm = 0;
for (const Keypoint& k : f.keypoints) {
min_s = std::min(min_s, (double)k.score);
max_s = std::max(max_s, (double)k.score);
}
for (size_t i = 0; i < f.keypoints.size(); ++i) {
double sq = 0;
for (int c = 0; c < f.desc_dim; ++c) {
const double v = f.descriptors[i * f.desc_dim + c];
sq += v * v;
}
mean_norm += std::sqrt(sq);
}
if (!f.keypoints.empty()) mean_norm /= (double)f.keypoints.size();
std::printf("%s %dx%d -> %zu keypoints in %.0f ms\n", image_path.c_str(), img.width,
img.height, f.keypoints.size(), ms);
std::printf(" score %.3f .. %.3f, mean |descriptor| %.5f\n", min_s, max_s, mean_norm);
// Descriptors are L2-normalized by construction; if they are not, nothing
// downstream that treats a dot product as a cosine is meaningful.
check(f.keypoints.empty() || std::fabs(mean_norm - 1.0) < 1e-3,
"descriptors are not unit norm (mean %.6f)", mean_norm);
check(min_s >= opts.min_score || f.keypoints.empty(),
"a keypoint scored %.4f, below min_score %.4f", min_s, opts.min_score);
if (out_path.empty()) return;
std::FILE* fp = std::fopen(out_path.c_str(), "wb");
check(fp != nullptr, "cannot write %s", out_path.c_str());
if (!fp) return;
const uint32_t version = 1, count = (uint32_t)f.keypoints.size(),
dim = (uint32_t)f.desc_dim;
std::fwrite("ALIKEDFT", 1, 8, fp);
std::fwrite(&version, 4, 1, fp);
std::fwrite(&f.width, 4, 1, fp);
std::fwrite(&f.height, 4, 1, fp);
std::fwrite(&count, 4, 1, fp);
std::fwrite(&dim, 4, 1, fp);
for (const Keypoint& k : f.keypoints) {
const float v[3] = {k.x, k.y, k.score};
std::fwrite(v, 4, 3, fp);
}
std::fwrite(f.descriptors.data(), 4, f.descriptors.size(), fp);
std::fclose(fp);
std::printf(" wrote %s\n", out_path.c_str());
}
// Match two dumps written by --image/--out, and append the correspondences to
// a third file so tools/aliked/compare_colmap.py can score them against ORT.
//
// u32 count, then count * { u32 i, u32 j, f32 score }
void run_matching(const std::string& model, const std::string& a_path,
const std::string& b_path, const std::string& out_path,
float min_score) {
struct Dump {
int w = 0, h = 0, dim = 0;
std::vector<float> xy, score, desc;
};
auto read = [&](const std::string& p) {
Dump d;
std::FILE* f = std::fopen(p.c_str(), "rb");
check(f != nullptr, "cannot read %s", p.c_str());
if (!f) return d;
char magic[8];
uint32_t version = 0, count = 0, dim = 0;
(void)std::fread(magic, 1, 8, f);
(void)std::fread(&version, 4, 1, f);
(void)std::fread(&d.w, 4, 1, f);
(void)std::fread(&d.h, 4, 1, f);
(void)std::fread(&count, 4, 1, f);
(void)std::fread(&dim, 4, 1, f);
d.dim = (int)dim;
d.xy.resize((size_t)count * 2);
d.score.resize(count);
for (uint32_t i = 0; i < count; i++) {
float v[3];
(void)std::fread(v, 4, 3, f);
d.xy[(size_t)i * 2] = v[0];
d.xy[(size_t)i * 2 + 1] = v[1];
d.score[i] = v[2];
}
d.desc.resize((size_t)count * dim);
(void)std::fread(d.desc.data(), 4, d.desc.size(), f);
std::fclose(f);
return d;
};
const Dump A = read(a_path), B = read(b_path);
if (A.desc.empty() || B.desc.empty()) return;
Matcher lg;
lg.load(model);
aliked::MatchInput ia, ib;
ia.keypoints = A.xy.data();
ia.descriptors = A.desc.data();
ia.n = (uint32_t)(A.xy.size() / 2);
ia.width = A.w;
ia.height = A.h;
ib.keypoints = B.xy.data();
ib.descriptors = B.desc.data();
ib.n = (uint32_t)(B.xy.size() / 2);
ib.width = B.w;
ib.height = B.h;
aliked::MatchOptions mo;
mo.min_score = min_score;
const double t0 = nn::now_ms();
std::vector<aliked::Match> m = lg.match(ia, ib, mo);
const double ms = nn::now_ms() - t0;
std::printf("lightglue: %u x %u -> %zu matches in %.0f ms\n", ia.n, ib.n, m.size(),
ms);
check(!m.empty(), "LightGlue returned no matches at all");
// Every index in range, every pair used once on each side: a mutual-nearest
// assignment cannot repeat either.
std::vector<char> seen_i(ia.n, 0), seen_j(ib.n, 0);
bool ok = true;
for (const aliked::Match& x : m) {
if (x.i >= ia.n || x.j >= ib.n) { ok = false; break; }
if (seen_i[x.i] || seen_j[x.j]) { ok = false; break; }
seen_i[x.i] = seen_j[x.j] = 1;
if (!(x.score >= min_score && x.score <= 1.0f + 1e-3f)) { ok = false; break; }
}
check(ok, "LightGlue produced a duplicate, out-of-range or invalid-score match");
if (out_path.empty()) return;
std::FILE* f = std::fopen(out_path.c_str(), "wb");
check(f != nullptr, "cannot write %s", out_path.c_str());
if (!f) return;
const uint32_t n = (uint32_t)m.size();
std::fwrite(&n, 4, 1, f);
for (const aliked::Match& x : m) {
std::fwrite(&x.i, 4, 1, f);
std::fwrite(&x.j, 4, 1, f);
std::fwrite(&x.score, 4, 1, f);
}
std::fclose(f);
std::printf(" wrote %s\n", out_path.c_str());
}
} // namespace
int main(int argc, char** argv) {
// Stage timings by default, but let $SSPLAT_NN_LOG win -- the tensor dumps
// in the model layer are how a mismatch against the reference is bisected,
// and forcing the level here would hide them.
if (!std::getenv("SSPLAT_NN_LOG")) nn::set_log_level(2);
bool fetch = false;
std::string image_path, out_path, model = "aliked-n16rot";
int max_image_size = 1600;
std::string match_a, match_b, lg_model = "aliked-lightglue";
ExtractOptions opts;
std::vector<std::string> explicit_paths;
for (int i = 1; i < argc; ++i) {
const std::string a = argv[i];
auto next = [&]() { return (i + 1 < argc) ? std::string(argv[++i]) : std::string(); };
if (a == "--fetch") fetch = true;
else if (a == "--image") image_path = next();
else if (a == "--out") out_path = next();
else if (a == "--model") model = next();
else if (a == "--max-image-size") max_image_size = std::atoi(next().c_str());
else if (a == "--max-features") opts.max_num_features = std::atoi(next().c_str());
else if (a == "--min-score") opts.min_score = (float)std::atof(next().c_str());
else if (a == "--match") { match_a = next(); match_b = next(); }
else if (a == "--lightglue-model") lg_model = next();
else explicit_paths.push_back(a);
}
try {
vk::Context::get();
std::vector<std::string> paths = explicit_paths;
if (paths.empty()) {
for (const char* id : {"aliked-n16rot", "aliked-n32"}) {
const ModelSource* src = find_model_source(id);
if (!src) continue;
if (fetch) {
paths.push_back(ensure_model(*src));
} else {
const std::string p = model_cache_path(*src);
std::error_code ec;
if (fs::exists(p, ec)) paths.push_back(p);
}
}
}
if (paths.empty()) {
std::printf("SKIP: no ALIKED checkpoint cached. Run with --fetch to "
"download from COLMAP's releases.\n");
vk::Stream::shutdown();
vk::Pipelines::get().shutdown();
vk::VramPool::get().releaseAll();
return 0;
}
for (const std::string& p : paths) check_checkpoint(p);
if (!match_a.empty()) {
std::printf("\nMatching\n");
run_matching(lg_model, match_a, match_b, out_path, 0.1f);
}
if (!image_path.empty()) {
std::printf("\nExtraction\n");
run_extraction(paths.empty() ? model : paths[0], image_path, out_path, opts,
max_image_size);
}
std::printf("\n%d checks, %d failures\n", g_checks, g_failures);
} catch (const std::exception& e) {
std::printf("EXCEPTION: %s\n", e.what());
return 1;
}
vk::Stream::shutdown();
vk::Pipelines::get().shutdown();
vk::VramPool::get().releaseAll();
std::printf("%s\n", g_failures == 0 ? "PASS" : "FAIL");
return g_failures == 0 ? 0 : 1;
}
src/app/cli/sfm_main.cpp
+30 -15
View File
@@ -157,7 +157,7 @@ static const CommandInfo kCommands[] = {
" 3 partial: under half the images registered, or over 2 px mean reprojection"},
{"extract", CMD_EXTRACT,
"detect SIFT features in an image or a directory of images",
"detect features in an image or a directory of images",
"<IMAGE|DIR> [-o OUT] [options]",
" GPU SIFT over one image or, recursively, a directory. A directory reuses one GPU\n"
" context and processes largest-first, so device buffers are allocated once;\n"
@@ -821,7 +821,7 @@ static int extractDirectory(const std::string& imagedir, const fs::path& outdir,
}
if (imgs.empty()) { fprintf(stderr, "no decodable images in %s\n", imagedir.c_str()); return 1; }
// Largest-first so SiftExtractor allocates device buffers exactly once.
// Largest-first so the extractor allocates device buffers exactly once.
std::vector<size_t> order(imgs.size());
for (size_t i = 0; i < order.size(); i++) order[i] = i;
std::stable_sort(order.begin(), order.end(), [&](size_t a, size_t b) {
@@ -888,11 +888,13 @@ static int extractDirectory(const std::string& imagedir, const fs::path& outdir,
((size_t)plan.num_threads * plan.decode_peak_bytes +
(size_t)plan.window * plan.held_bytes) >> 20);
SiftExtractor ext(opt);
std::unique_ptr<IFeatureExtractor> ext =
createFeatureExtractor(cfg.features, opt, cfg.aliked);
if (opt.verbose) fprintf(stderr, "[extract] frontend: %s\n", ext->name());
loadImagesInOrder(
paths, plan, lopt,
[&](size_t k, GrayImage& img) {
FeatureSet f = ext.extract(img);
FeatureSet f = ext->extract(img);
sampleFeatureColors(f, img);
uint32_t dropped = 0;
if (!lopt.mask_paths.empty() && !lopt.mask_paths[k].empty()) {
@@ -1005,8 +1007,9 @@ static int cmdExtract(int argc, char** argv) {
GrayImage img = loadGrayImage(image, cfg.max_image_size, /*want_color=*/true, maskpath);
if (cfg.sift.verbose)
fprintf(stderr, "[sfm] %s -> %dx%d gray\n", image.c_str(), img.width, img.height);
SiftExtractor ext(cfg.sift);
FeatureSet fset = ext.extract(img);
std::unique_ptr<IFeatureExtractor> ext =
createFeatureExtractor(cfg.features, cfg.sift, cfg.aliked);
FeatureSet fset = ext->extract(img);
sampleFeatureColors(fset, img);
uint32_t masked_out = applyMask(fset, img.mask);
finishFeatures(fset, img);
@@ -1139,9 +1142,12 @@ static int matchFeatureDir(const std::string& featdir, const SfmConfig& cfg, Pai
fprintf(stderr, "[match] pair selection: top-%u features, %u neighbors\n",
popt.num_features, popt.num_neighbors);
BruteForceMatcher matcher(opt);
std::unique_ptr<IFeatureMatcher> matcher =
createFeatureMatcher(cfg.matcher, opt, cfg.lightglue);
if (verbose && cfg.matcher != "bruteforce")
fprintf(stderr, "[match] matcher: %s\n", matcher->name());
auto matchFn = [&](size_t b, size_t e, std::vector<std::vector<FeatureMatch>>& mout) {
matcher.matchBatch(feats, pairs, b, e, mout);
matcher->matchBatch(feats, pairs, b, e, mout);
};
std::function<void(size_t, size_t)> progress;
if (verbose)
@@ -1183,7 +1189,7 @@ static int matchFeatureDir(const std::string& featdir, const SfmConfig& cfg, Pai
std::vector<std::vector<FeatureMatch>> chunk;
for (size_t b = 0; b < sample.size(); b += 16) {
size_t e = std::min(b + 16, sample.size());
matcher.matchBatch(feats, sample, b, e, chunk);
matcher->matchBatch(feats, sample, b, e, chunk);
for (size_t k = b; k < e; k++) sm.push_back(std::move(chunk[k - b]));
}
}
@@ -1934,12 +1940,21 @@ int ssplat_sfm_main(int argc, char** argv) {
std::printf("%s %s\n", kProgram, SSPLAT_VERSION);
return 0;
}
if (cmd == "auto") return cmdAuto(argc - 2, argv + 2);
if (cmd == "extract") return cmdExtract(argc - 2, argv + 2);
if (cmd == "match") return cmdMatch(argc - 2, argv + 2);
if (cmd == "map") return cmdMap(argc - 2, argv + 2);
if (cmd == "merge") return cmdMerge(argc - 2, argv + 2);
if (cmd == "ba") return cmdBa(argc - 2, argv + 2);
// One catch for every subcommand. Setup failures throw rather than return
// -- a checkpoint that will not download, a matcher handed the wrong kind
// of descriptor -- and those messages are written to be read by the person
// who typed the command, not by a terminate handler.
try {
if (cmd == "auto") return cmdAuto(argc - 2, argv + 2);
if (cmd == "extract") return cmdExtract(argc - 2, argv + 2);
if (cmd == "match") return cmdMatch(argc - 2, argv + 2);
if (cmd == "map") return cmdMap(argc - 2, argv + 2);
if (cmd == "merge") return cmdMerge(argc - 2, argv + 2);
if (cmd == "ba") return cmdBa(argc - 2, argv + 2);
} catch (const std::exception& e) {
std::fprintf(stderr, "%s %s: error: %s\n", kProgram, cmd.c_str(), e.what());
return 1;
}
std::fprintf(stderr, "%s: error: unknown command '%s'\n", kProgram, cmd.c_str());
std::fprintf(stderr, "Try '%s --help' for the list of commands.\n", kProgram);
return 1;
src/app/gui/GuiApp.cpp
+38 -2
View File
@@ -1050,6 +1050,39 @@ void GuiApp::draw_sfm_advanced() {
"What the input is, which sets the pairing strategy and how forgiving "
"the mapper is. Set from the input type when you picked it.");
const char* frontends[] = {"SIFT (classic)", "ALIKED N16-rot (learned)",
"ALIKED N32 (learned, wider)"};
ImGui::SetNextItemWidth(260);
ImGui::Combo("Features", &_sfm_job.features, frontends, 3);
help_tooltip_on_hover(
"Which detector and descriptor. SIFT is the classic one and needs "
"nothing downloaded. The ALIKED options are a learned frontend: they "
"fetch a small checkpoint (3-4 MB) on first use, find fewer but "
"better-localized keypoints, and match markedly more image pairs on "
"hard captures. N32 samples more positions per descriptor -- slower, "
"slightly stronger.");
{
// Brute force is the only option for SIFT, so say so by disabling the
// combo rather than by letting the run fail.
const bool learned = _sfm_job.features != 0;
const char* matchers[] = {"Brute force", "LightGlue (learned)"};
ImGui::BeginDisabled(!learned);
ImGui::SetNextItemWidth(260);
int shown = learned ? _sfm_job.matcher : 0;
if (ImGui::Combo("Matcher", &shown, matchers, 2) && learned)
_sfm_job.matcher = shown;
ImGui::EndDisabled();
help_tooltip_on_hover(
learned
? "How descriptors are matched. LightGlue is a learned matcher: "
"it finds far more correct correspondences on hard pairs, and "
"costs tens of milliseconds per pair instead of a few, so it "
"runs behind pair selection."
: "LightGlue needs the learned descriptors -- pick an ALIKED "
"frontend above to enable it.");
}
const char* mappers[] = {"Automatic", "Flat (one reconstruction)",
"Bottom-up (atoms, merged upwards)"};
ImGui::SetNextItemWidth(260);
@@ -1079,8 +1112,11 @@ void GuiApp::draw_sfm_advanced() {
ImGui::SetNextItemWidth(260);
ImGui::InputInt("Max features per image (0 = auto)",
&_sfm_job.max_features, 0, 0);
help_tooltip_on_hover("Keypoints kept per image, largest scales first. "
"Overrides the quality preset when non-zero.");
help_tooltip_on_hover("Keypoints kept per image -- largest scales first "
"for SIFT, highest detection scores for a learned "
"frontend. Overrides the quality preset when "
"non-zero. The two are not comparable: SIFT wants "
"tens of thousands, ALIKED a few thousand.");
ImGui::SetNextItemWidth(260);
ImGui::InputInt("Max image size (0 = auto)", &_sfm_job.max_image_size, 0, 0);
help_tooltip_on_hover(
src/app/gui/SfmRunner.cpp
+13 -1
View File
@@ -23,6 +23,8 @@ const char* kDataType[] = {"individual", "video", "internet"};
const char* kCameraMode[] = {"single", "folder", "image"};
const char* kPairs[] = {"auto", "exhaustive", "sequential", "prefilter"};
const char* kMapper[] = {"auto", "flat", "bottom-up"};
const char* kFeatures[] = {"sift", "aliked-n16rot", "aliked-n32"};
const char* kMatcher[] = {"bruteforce", "lightglue"};
template <int N>
const char* pick(const char* const (&table)[N], int i, int fallback = 0) {
@@ -234,6 +236,11 @@ void SfmRunner::run(SfmJob job) {
"--camera-model", job.camera_model,
"--camera-mode", pick(kCameraMode, job.camera_mode, 1),
"--mapper", pick(kMapper, job.mapper),
"--features", pick(kFeatures, job.features),
// LightGlue only exists for the learned descriptors; asking
// for it with SIFT selected is a usage error, so do not.
"--matcher",
pick(kMatcher, job.features == 0 ? 0 : job.matcher),
};
if (job.pairs > 0) {
argv.push_back("--pairs");
@@ -250,7 +257,12 @@ void SfmRunner::run(SfmJob job) {
argv.push_back(buf);
}
if (job.max_features > 0) {
argv.push_back("--max-features");
// Each frontend has its own count flag, because their budgets
// are not comparable -- a learned detector emits a few
// thousand better-localized points where SIFT wants tens of
// thousands. One spinner, routed to whichever is running.
argv.push_back(job.features == 0 ? "--max-features"
: "--aliked-max-features");
argv.push_back(std::to_string(job.max_features));
}
if (job.max_image_size > 0) {
src/app/gui/SfmRunner.h
+9
View File
@@ -69,6 +69,15 @@ struct SfmJob {
int max_features = 0; // 0 = the quality preset's
int max_image_size = 0; // 0 = the quality preset's
int mapper = 0; // 0 auto, 1 flat, 2 bottom-up
// 0 SIFT, 1 ALIKED-n16rot, 2 ALIKED-n32. The learned ones fetch a
// checkpoint on first use and run on their own resolution ladder, so the
// quality preset means something different for each -- which is why this
// is a frontend choice and not a quality level.
int features = 0;
// 0 brute force, 1 LightGlue. Only meaningful with a learned frontend, and
// an order of magnitude slower per pair -- the panel greys it out for SIFT
// and the CLI refuses the combination outright.
int matcher = 0;
bool keep_intermediate = false; // keep features/ and matches.bin
// Extra flags typed by the user, appended verbatim. The escape hatch for
src/nn/OpConv.cpp
+104
View File
@@ -18,6 +18,20 @@ struct Im2ColParams {
uint32_t p0, P, groups_per_row;
};
struct DeformIm2ColParams {
uint64_t out, x, offset;
uint32_t Hi, Wi, Ci, Ho, Wo, kh, kw;
uint32_t stride_y, stride_x, pad_y, pad_x;
uint32_t p0, P;
float max_offset;
uint32_t groups_per_row;
};
struct PatchGatherParams {
uint64_t out, x, centers;
uint32_t Hi, Wi, C, N, k, groups_per_row;
};
struct DepthwiseParams {
uint64_t out, x, w, bias;
uint32_t Hi, Wi, Ho, Wo, C, kh, kw;
@@ -102,6 +116,96 @@ void conv2d(vk::Arena& arena, const Tensor& out, const Tensor& in, const Tensor&
}
}
void deform_conv2d(vk::Arena& arena, const Tensor& out, const Tensor& in,
const Tensor& offset, const Tensor& w_in, int kh, int kw,
float max_offset, const ConvOpts& o) {
NN_CHECK(out.ndim == 3 && in.ndim == 3 && offset.ndim == 3,
"deform_conv2d expects [H, W, C] tensors");
const Tensor w = w_in.asMatrix();
const int64_t Hi = in.shape[0], Wi = in.shape[1], Ci = in.shape[2];
const int64_t Ho = out.shape[0], Wo = out.shape[1], Co = out.shape[2];
const int64_t K = Ci * kh * kw;
NN_CHECK(w.rows() == Co && w.cols() == K,
"deform_conv2d: weight is [%lld, %lld] but %lldx%lldx%lld -> %lld needs "
"[%lld, %lld]",
(long long)w.rows(), (long long)w.cols(), (long long)kh, (long long)kw,
(long long)Ci, (long long)Co, (long long)Co, (long long)K);
NN_CHECK(offset.shape[0] == Ho && offset.shape[1] == Wo &&
offset.shape[2] == 2 * kh * kw,
"deform_conv2d: offset is [%lld, %lld, %lld], expected [%lld, %lld, %lld]",
(long long)offset.shape[0], (long long)offset.shape[1],
(long long)offset.shape[2], (long long)Ho, (long long)Wo,
(long long)(2 * kh * kw));
// The offset map is read as raw f32 by the kernel; an f16 one would need a
// second load path for two values per tap and buys nothing at this size.
NN_CHECK(offset.dtype == DType::F32, "deform_conv2d: offset must be f32");
const int64_t positions = Ho * Wo;
int64_t chunk = std::max<int64_t>(64, kMaxColBytes / (K * 4));
chunk = std::min(chunk, positions);
vk::ArenaScope scope(arena);
Tensor cols = arena_tensor(arena, DType::F32, chunk, K);
for (int64_t p0 = 0; p0 < positions; p0 += chunk) {
const int64_t P = std::min(chunk, positions - p0);
DeformIm2ColParams ip{};
ip.out = cols.ptr;
ip.x = in.ptr;
ip.offset = offset.ptr;
ip.Hi = (uint32_t)Hi;
ip.Wi = (uint32_t)Wi;
ip.Ci = (uint32_t)Ci;
ip.Ho = (uint32_t)Ho;
ip.Wo = (uint32_t)Wo;
ip.kh = (uint32_t)kh;
ip.kw = (uint32_t)kw;
ip.stride_y = (uint32_t)o.stride_y;
ip.stride_x = (uint32_t)o.stride_x;
ip.pad_y = (uint32_t)o.pad_y;
ip.pad_x = (uint32_t)o.pad_x;
ip.p0 = (uint32_t)p0;
ip.P = (uint32_t)P;
ip.max_offset = max_offset;
vk::SpecList spec{(uint32_t)(in.dtype == DType::F16), 0u};
vk::Stream::get().dispatchFlat("conv.deform_im2col", spec, P * K, 256, &ip,
sizeof(ip), &ip.groups_per_row);
LinearOpts lo;
lo.bias = o.bias;
lo.act = o.act;
Tensor out_chunk(out.ptr + (uint64_t)(p0 * Co) * 4, DType::F32, P, Co);
linear(out_chunk, cols.view(P, K), w, lo);
}
}
void patch_gather(const Tensor& out, const Tensor& in, const Tensor& centers, int k) {
NN_CHECK(in.ndim == 3, "patch_gather expects an [H, W, C] map");
NN_CHECK(centers.dtype == DType::I32 && centers.cols() == 2,
"patch_gather: centers must be an i32 [N, 2] tensor");
const int64_t N = centers.rows();
const int64_t C = in.shape[2];
NN_CHECK(out.rows() == N && out.cols() == C * k * k,
"patch_gather: out is [%lld, %lld], expected [%lld, %lld]",
(long long)out.rows(), (long long)out.cols(), (long long)N,
(long long)(C * k * k));
if (N == 0) return;
PatchGatherParams p{};
p.out = out.ptr;
p.x = in.ptr;
p.centers = centers.ptr;
p.Hi = (uint32_t)in.shape[0];
p.Wi = (uint32_t)in.shape[1];
p.C = (uint32_t)C;
p.N = (uint32_t)N;
p.k = (uint32_t)k;
vk::SpecList spec{(uint32_t)(in.dtype == DType::F16), 0u};
vk::Stream::get().dispatchFlat("conv.patch_gather", spec, N * C * k * k, 256, &p,
sizeof(p), &p.groups_per_row);
}
void conv2d_depthwise(const Tensor& out, const Tensor& in, const Tensor& w_in, int kh,
int kw, const ConvOpts& o) {
NN_CHECK(out.ndim == 3 && in.ndim == 3, "conv2d_depthwise expects [H, W, C]");
src/nn/OpMisc.cpp
+96 -4
View File
@@ -66,6 +66,23 @@ struct MaskExportParams {
uint32_t groups_per_row;
};
struct PoolParams {
uint64_t out, x;
uint32_t Ho, Wo, Hi, Wi, C, kh, kw, stride_y, stride_x, groups_per_row;
};
struct GridSampleParams {
uint64_t out, x, pos;
uint32_t H, W, C, N, groups_per_row;
};
struct NormalizeParams {
uint64_t out, x;
uint32_t rows, cols;
float eps;
uint32_t groups_per_row;
};
} // namespace
// ================
@@ -250,7 +267,8 @@ void roi_align(const Tensor& out, const Tensor& feat, const Tensor& boxes, int H
// Resample
// ================
static void resize_op(const char* entry, const Tensor& out, const Tensor& in) {
static void resize_op(const char* entry, const Tensor& out, const Tensor& in,
bool align_corners = false) {
NN_CHECK(out.ndim == 3 && in.ndim == 3,
"%s expects [H, W, C] tensors (got %dD and %dD)", entry, out.ndim,
in.ndim);
@@ -264,13 +282,13 @@ static void resize_op(const char* entry, const Tensor& out, const Tensor& in) {
p.Hi = (uint32_t)in.shape[0];
p.Wi = (uint32_t)in.shape[1];
p.C = (uint32_t)out.shape[2];
vk::SpecList spec{(uint32_t)(in.dtype == DType::F16), 0u};
vk::SpecList spec{(uint32_t)(in.dtype == DType::F16), 0u, (uint32_t)align_corners};
vk::Stream::get().dispatchFlat(entry, spec, out.numel(), 256, &p, sizeof(p),
&p.groups_per_row);
}
void resize_bilinear(const Tensor& out, const Tensor& in) {
resize_op("resample.resize_bilinear", out, in);
void resize_bilinear(const Tensor& out, const Tensor& in, bool align_corners) {
resize_op("resample.resize_bilinear", out, in, align_corners);
}
void upsample_nearest2x(const Tensor& out, const Tensor& in) {
resize_op("resample.upsample_nearest2x", out, in);
@@ -279,6 +297,80 @@ void maxpool2x2(const Tensor& out, const Tensor& in) {
resize_op("resample.maxpool2x2", out, in);
}
void avgpool(const Tensor& out, const Tensor& in, int kernel, int stride) {
NN_CHECK(out.ndim == 3 && in.ndim == 3, "avgpool expects [H, W, C] tensors");
NN_CHECK(out.shape[2] == in.shape[2], "avgpool: channel counts differ");
NN_CHECK(kernel > 0, "avgpool: kernel must be positive");
if (stride <= 0) stride = kernel;
// torch with ceil_mode=False; say so here rather than let a mis-sized
// output silently read past the last complete window.
const int64_t want_h = (in.shape[0] - kernel) / stride + 1;
const int64_t want_w = (in.shape[1] - kernel) / stride + 1;
NN_CHECK(out.shape[0] == want_h && out.shape[1] == want_w,
"avgpool(%d, %d) of %lldx%lld is %lldx%lld, but out is %lldx%lld", kernel,
stride, (long long)in.shape[0], (long long)in.shape[1], (long long)want_h,
(long long)want_w, (long long)out.shape[0], (long long)out.shape[1]);
PoolParams p{};
p.out = out.ptr;
p.x = in.ptr;
p.Ho = (uint32_t)out.shape[0];
p.Wo = (uint32_t)out.shape[1];
p.Hi = (uint32_t)in.shape[0];
p.Wi = (uint32_t)in.shape[1];
p.C = (uint32_t)out.shape[2];
p.kh = p.kw = (uint32_t)kernel;
p.stride_y = p.stride_x = (uint32_t)stride;
vk::SpecList spec{(uint32_t)(in.dtype == DType::F16), 0u, 0u};
vk::Stream::get().dispatchFlat("resample.avgpool", spec, out.numel(), 256, &p,
sizeof(p), &p.groups_per_row);
}
void grid_sample_points(const Tensor& out, const Tensor& in, const Tensor& pos,
bool align_corners) {
NN_CHECK(in.ndim == 3, "grid_sample_points expects an [H, W, C] map");
NN_CHECK(pos.dtype == DType::F32 && pos.cols() == 2,
"grid_sample_points: pos must be an f32 [N, 2] tensor");
const int64_t N = pos.rows(), C = in.shape[2];
NN_CHECK(out.rows() == N && out.cols() == C,
"grid_sample_points: out is [%lld, %lld], expected [%lld, %lld]",
(long long)out.rows(), (long long)out.cols(), (long long)N, (long long)C);
if (N == 0) return;
GridSampleParams p{};
p.out = out.ptr;
p.x = in.ptr;
p.pos = pos.ptr;
p.H = (uint32_t)in.shape[0];
p.W = (uint32_t)in.shape[1];
p.C = (uint32_t)C;
p.N = (uint32_t)N;
vk::SpecList spec{(uint32_t)(in.dtype == DType::F16), 0u, (uint32_t)align_corners};
vk::Stream::get().dispatchFlat("misc.grid_sample_points", spec, N * C, 256, &p,
sizeof(p), &p.groups_per_row);
}
void l2_normalize_rows(const Tensor& out, const Tensor& x, float eps) {
NN_CHECK(out.rows() == x.rows() && out.cols() == x.cols(),
"l2_normalize_rows: shapes differ");
const int64_t rows = x.rows();
if (rows == 0) return;
NormalizeParams p{};
p.out = out.ptr;
p.x = x.ptr;
p.rows = (uint32_t)rows;
p.cols = (uint32_t)x.cols();
p.eps = eps;
// One workgroup per row, not one thread per element: the grid folds over
// rows, and the shader reconstructs the row from (gid.y, gid.x).
const vk::Stream::Fold fold = vk::Stream::fold1D(rows, 1);
p.groups_per_row = fold.per_row;
vk::SpecList spec{(uint32_t)(x.dtype == DType::F16), 0u, 0u};
vk::Stream::get().dispatch("misc.l2_normalize_rows", spec, fold.per_row, fold.rows, 1,
&p, sizeof(p));
}
void resize_binarize(const Tensor& out_u8, const Tensor& logits, int64_t Ho, int64_t Wo,
float threshold) {
NN_CHECK(out_u8.dtype == DType::U8, "resize_binarize writes a u8 tensor");
src/nn/Ops.h
+38 -1
View File
@@ -26,6 +26,8 @@ enum class Act : uint32_t {
GeluErf = 2, // nn.GELU() -- SAM 3's ViT/neck/decoder use this exact form
GeluTanh = 3,
Sigmoid = 4,
Selu = 5, // nn.SELU() -- ALIKED's gate, in every block and both heads
LogSigmoid = 6,// F.logsigmoid -- LightGlue's matchability term
};
enum class AttnBias : uint32_t {
@@ -155,6 +157,25 @@ void conv2d(vk::Arena& arena, const Tensor& out, const Tensor& in, const Tensor&
void conv2d_depthwise(const Tensor& out, const Tensor& in, const Tensor& w, int kh,
int kw, const ConvOpts& opts = {});
// torchvision's deform_conv2d with mask = None, groups = 1.
//
// `offset` is [Ho, Wo, 2*kh*kw] -- channel-last, (dy, dx) per tap in kernel
// order, which is what a plain conv2d producing 2*kh*kw channels lands in with
// no permute. `max_offset > 0` clamps each component to +-max_offset before
// sampling (ALIKED clamps to max(H, W) / 4); pass 0 for no clamp.
//
// Runs as chunked deform-im2col + the same GEMM and the same weight layout as
// conv2d, so a deformable conv costs one address computation more per tap than
// a normal one.
void deform_conv2d(vk::Arena& arena, const Tensor& out, const Tensor& in,
const Tensor& offset, const Tensor& w, int kh, int kw,
float max_offset = 0.0f, const ConvOpts& opts = {});
// out[N, C*k*k] = the k x k patch of `in` centred on each of N integer
// (x, y) centres, in the column order conv2d's weight expects. Out-of-range
// taps read zero. `centers` is an I32 [N, 2] tensor.
void patch_gather(const Tensor& out, const Tensor& in, const Tensor& centers, int k);
// ConvTranspose2d(kernel=2, stride=2). `w_packed` is the checkpoint weight
// repacked to [Cout*4, Cin] at load time (see model/Weights.cpp); the four
// kernel taps become four output-channel groups, so the tuned GEMM does the
@@ -166,10 +187,26 @@ void conv_transpose2x2(vk::Arena& arena, const Tensor& out, const Tensor& in,
// out[(H/p)*(W/p), p*p*C] with column order c*p*p + ky*p + kx.
void patchify(const Tensor& out, const Tensor& in, int patch);
void resize_bilinear(const Tensor& out, const Tensor& in);
// `align_corners` picks between torch's two mappings and is NOT cosmetic:
// false (the default here, and what mask upsampling depends on) maps
// src = (dst + 0.5) * scale - 0.5; true maps src = dst * (Hi-1)/(Ho-1).
void resize_bilinear(const Tensor& out, const Tensor& in, bool align_corners = false);
void upsample_nearest2x(const Tensor& out, const Tensor& in);
void maxpool2x2(const Tensor& out, const Tensor& in);
// nn.AvgPool2d(kernel, stride), no padding, ceil_mode = False. `out` must be
// sized [(Hi-kernel)/stride + 1, ...]; stride 0 means "same as kernel".
void avgpool(const Tensor& out, const Tensor& in, int kernel, int stride = 0);
// out[N, C] = bilinear sample of in[H, W, C] at `pos`, an f32 [N, 2] tensor of
// normalized (x, y) in [-1, 1]. Reads zero outside, i.e. torch's
// padding_mode='zeros'. See the shader for why align_corners matters.
void grid_sample_points(const Tensor& out, const Tensor& in, const Tensor& pos,
bool align_corners = true);
// F.normalize(x, p=2, dim=-1). `out` may alias `x`.
void l2_normalize_rows(const Tensor& out, const Tensor& x, float eps = 1e-12f);
// Resize a single-channel logit map to [Ho, Wo], threshold, and write packed
// 0/255 bytes. `out` is a U8 tensor; the buffer must be 4-byte rounded (every
// allocation is).
src/nn/shaders/_common.slang
+40
View File
@@ -65,6 +65,13 @@ static const int ACT_RELU = 1;
static const int ACT_GELU_ERF = 2; // nn.GELU() -- the PyTorch default
static const int ACT_GELU_TANH = 3; // nn.GELU('tanh')
static const int ACT_SIGMOID = 4;
static const int ACT_SELU = 5; // nn.SELU() -- ALIKED's gate
static const int ACT_LOGSIGMOID = 6; // F.logsigmoid -- LightGlue
// nn.SELU's constants, which are fixed by the paper (Klambauer et al. 2017)
// and are what torch hard-codes -- not tunable, and not the same as ELU's.
static const float SELU_ALPHA = 1.6732632423543772848170429916717f;
static const float SELU_SCALE = 1.0507009873554804934193349852946f;
// erf-exact GELU: 0.5x(1 + erf(x/sqrt2)).
//
@@ -95,6 +102,10 @@ float apply_act(float x, int kind) {
case ACT_GELU_ERF: return gelu_erf(x);
case ACT_GELU_TANH: return gelu_tanh(x);
case ACT_SIGMOID: return 1.0f / (1.0f + exp(-x));
case ACT_SELU: return SELU_SCALE * ((x > 0.0f) ? x : SELU_ALPHA * (exp(x) - 1.0f));
// log(sigmoid(x)), in the form that does not overflow either tail:
// min(x,0) - log(1 + exp(-|x|)).
case ACT_LOGSIGMOID: return min(x, 0.0f) - log(1.0f + exp(-abs(x)));
default: return x;
}
}
@@ -140,4 +151,33 @@ float sample_bilinear(uint* src, bool src_f16, int H, int W, int C, int c, float
fy * ((1.0f - fx) * v10 + fx * v11);
}
// Bilinear fetch that reads ZERO outside the map instead of clamping to the
// edge -- torchvision's rule for deformable convolution and torch's
// `grid_sample(padding_mode='zeros')`.
//
// The corner-by-corner guards matter and are not the same as an early
// bounds test on (y, x): at y = H - 0.5 the upper row is out of range but
// contributes weight 0.5, and torchvision takes that contribution as zero
// rather than repeating row H-1. Getting this wrong is invisible in the
// interior and shifts every response near the border.
float sample_bilinear_zero(uint* src, bool src_f16, int H, int W, int C, int c, float y,
float x) {
if (y <= -1.0f || float(H) <= y || x <= -1.0f || float(W) <= x) return 0.0f;
int y0 = int(floor(y));
int x0 = int(floor(x));
int y1 = y0 + 1;
int x1 = x0 + 1;
float ly = y - float(y0);
float lx = x - float(x0);
float hy = 1.0f - ly;
float hx = 1.0f - lx;
uint stride_y = uint(W) * uint(C);
float v00 = 0.0f, v01 = 0.0f, v10 = 0.0f, v11 = 0.0f;
if (y0 >= 0 && x0 >= 0) v00 = load_elem(src, uint(y0) * stride_y + uint(x0) * uint(C) + uint(c), src_f16);
if (y0 >= 0 && x1 <= W - 1) v01 = load_elem(src, uint(y0) * stride_y + uint(x1) * uint(C) + uint(c), src_f16);
if (y1 <= H - 1 && x0 >= 0) v10 = load_elem(src, uint(y1) * stride_y + uint(x0) * uint(C) + uint(c), src_f16);
if (y1 <= H - 1 && x1 <= W - 1) v11 = load_elem(src, uint(y1) * stride_y + uint(x1) * uint(C) + uint(c), src_f16);
return hy * (hx * v00 + lx * v01) + ly * (hx * v10 + lx * v11);
}
#endif // NN_COMMON_INCLUDED
src/nn/shaders/conv.slang
+131
View File
@@ -70,6 +70,137 @@ void im2col(uint3 gid : SV_GroupID, uint3 tid : SV_GroupThreadID,
p.out[i] = v;
}
// ---------------------------------------------------------------------------
// Deformable convolution
// ---------------------------------------------------------------------------
struct DeformIm2ColParams {
float* out; // [P, Ci*kh*kw]
uint* x; // [Hi, Wi, Ci]
uint* offset; // [Ho, Wo, 2*kh*kw] -- (dy, dx) per tap, channel-last
uint Hi;
uint Wi;
uint Ci;
uint Ho;
uint Wo;
uint kh;
uint kw;
uint stride_y;
uint stride_x;
uint pad_y;
uint pad_x;
uint p0;
uint P;
float max_offset; // offsets are clamped to +-this; <= 0 disables
uint groups_per_row;
};
// The same column layout im2col produces, so the GEMM and the checkpoint
// weight are shared verbatim between the two -- only where a tap reads
// changes.
//
// This is torchvision's `deform_conv2d` with mask = None, which is the form
// ALIKED's checkpoints carry (their offset convs have 2*k*k output channels,
// not 3*k*k; AlikedWeights::load refuses anything else). Two details are
// load-bearing and neither is guessable from the name:
//
// * The sampling position is p0 + tap + offset, where the offset for tap
// (ky, kx) lives at offset channels 2*(ky*kw + kx) and +1, in that order,
// (dy, dx). Swapping them produces a network that trains fine and
// matches badly.
// * Sampling reads ZERO outside the input, not the edge pixel. See
// sample_bilinear_zero.
//
// One bilinear fetch per (position, tap, channel): the four corner loads are
// what a plain im2col would do anyway, and only the weights are recomputed
// per channel.
[shader("compute")]
[numthreads(WG, 1, 1)]
void deform_im2col(uint3 gid : SV_GroupID, uint3 tid : SV_GroupThreadID,
uniform DeformIm2ColParams p) {
uint i = fold_index(gid, tid, p.groups_per_row, uint(WG));
uint K = p.Ci * p.kh * p.kw;
if (i >= p.P * K) return;
uint col = i % K;
uint pos = p.p0 + i / K;
uint kx = col % p.kw;
uint rest = col / p.kw;
uint ky = rest % p.kh;
uint ci = rest / p.kh;
uint xo = pos % p.Wo;
uint yo = pos / p.Wo;
uint tap = ky * p.kw + kx;
uint obase = (yo * p.Wo + xo) * (2u * p.kh * p.kw) + 2u * tap;
float dy = asfloat(p.offset[obase]);
float dx = asfloat(p.offset[obase + 1u]);
if (p.max_offset > 0.0f) {
dy = clamp(dy, -p.max_offset, p.max_offset);
dx = clamp(dx, -p.max_offset, p.max_offset);
}
float y = float(int(yo * p.stride_y + ky) - int(p.pad_y)) + dy;
float x = float(int(xo * p.stride_x + kx) - int(p.pad_x)) + dx;
p.out[i] = sample_bilinear_zero(p.x, kInputF16 != 0, int(p.Hi), int(p.Wi), int(p.Ci),
int(ci), y, x);
}
// ---------------------------------------------------------------------------
// Patch gather
// ---------------------------------------------------------------------------
struct PatchGatherParams {
float* out; // [N, C*k*k]
uint* x; // [Hi, Wi, C]
uint* centers; // [N, 2] i32 (x, y) pixel coordinates
uint Hi;
uint Wi;
uint C;
uint N;
uint k;
uint groups_per_row;
};
// A k x k patch around each of N integer centers, laid out so the result
// multiplies a [Cout, C*k*k] conv weight verbatim -- im2col at arbitrary
// positions rather than at every position.
//
// What ALIKED's descriptor head does first: its offset network is a
// Conv2d(dim, 2M, k=3, padding=0) evaluated at the keypoints only, which is
// this gather and then one GEMM, instead of a convolution over a whole image
// to read a few thousand pixels of.
//
// Out-of-range taps read zero, matching a VALID convolution of a
// zero-padded map.
[shader("compute")]
[numthreads(WG, 1, 1)]
void patch_gather(uint3 gid : SV_GroupID, uint3 tid : SV_GroupThreadID,
uniform PatchGatherParams p) {
uint i = fold_index(gid, tid, p.groups_per_row, uint(WG));
uint K = p.C * p.k * p.k;
if (i >= p.N * K) return;
uint col = i % K;
uint n = i / K;
uint kx = col % p.k;
uint rest = col / p.k;
uint ky = rest % p.k;
uint ci = rest / p.k;
int cx = int(p.centers[2u * n]);
int cy = int(p.centers[2u * n + 1u]);
int r = int(p.k) / 2;
int y = cy - r + int(ky);
int x = cx - r + int(kx);
float v = 0.0f;
if (y >= 0 && y < int(p.Hi) && x >= 0 && x < int(p.Wi))
v = load_elem(p.x, (uint(y) * p.Wi + uint(x)) * p.C + ci, kInputF16 != 0);
p.out[i] = v;
}
struct DepthwiseParams {
float* out; // [Ho, Wo, C]
uint* x; // [Hi, Wi, C]
src/nn/shaders/misc.slang
+102 -1
View File
@@ -5,9 +5,10 @@
#include "_common.slang"
// Spec ids 0..1 (module-wide).
// Spec ids 0..2 (module-wide).
[SpecializationConstant] const int kInputF16 = 0;
[SpecializationConstant] const int kAct = 0;
[SpecializationConstant] const int kAlignCorners = 0;
static const int WG = 256;
@@ -308,3 +309,103 @@ void roi_align(uint3 gid : SV_GroupID, uint3 tid : SV_GroupThreadID,
}
p.out[i] = sum / float(ny * nx);
}
// ---------------------------------------------------------------------------
// Point-wise grid sample
// ---------------------------------------------------------------------------
struct GridSampleParams {
float* out; // [N, C]
uint* x; // [H, W, C]
uint* pos; // [N, 2] f32, normalized (gx, gy) in [-1, 1]
uint H;
uint W;
uint C;
uint N;
uint groups_per_row;
};
// torch's `grid_sample(mode='bilinear', padding_mode='zeros')` at N arbitrary
// points, with align_corners taken from the module's spec constant.
//
// A whole-map resample (resample.slang) is the wrong shape for this: ALIKED's
// descriptor head samples M positions around each of a few thousand keypoints,
// i.e. a list, not a grid. roi_align is the closest existing op and is not a
// substitute -- it averages over a box, this reads one point.
//
// The two align_corners conventions differ in what -1 and +1 mean: the corner
// *pixel centres* when true, the outer *edges* of the corner pixels when
// false. ALIKED normalizes its keypoints as x / (W - 1) * 2 - 1, which is the
// true convention, and mixing them costs half a pixel at the border.
[shader("compute")]
[numthreads(WG, 1, 1)]
void grid_sample_points(uint3 gid : SV_GroupID, uint3 tid : SV_GroupThreadID,
uniform GridSampleParams p) {
uint i = fold_index(gid, tid, p.groups_per_row, uint(WG));
if (i >= p.N * p.C) return;
uint c = i % p.C;
uint n = i / p.C;
float gx = asfloat(p.pos[2u * n]);
float gy = asfloat(p.pos[2u * n + 1u]);
float sx, sy;
if (kAlignCorners != 0) {
sx = (gx + 1.0f) * 0.5f * (float(p.W) - 1.0f);
sy = (gy + 1.0f) * 0.5f * (float(p.H) - 1.0f);
} else {
sx = ((gx + 1.0f) * float(p.W) - 1.0f) * 0.5f;
sy = ((gy + 1.0f) * float(p.H) - 1.0f) * 0.5f;
}
p.out[i] = sample_bilinear_zero(p.x, kInputF16 != 0, int(p.H), int(p.W), int(p.C),
int(c), sy, sx);
}
// ---------------------------------------------------------------------------
// Row-wise L2 normalization
// ---------------------------------------------------------------------------
struct NormalizeParams {
float* out; // [rows, cols]
uint* x;
uint rows;
uint cols;
float eps;
uint groups_per_row;
};
// out[r] = x[r] / max(||x[r]||, eps) -- torch's F.normalize(p=2, dim=-1).
//
// One workgroup per row with a groupshared tree reduction, deliberately not
// wave intrinsics: a descriptor that differs in its last bits between a wave32
// and a wave64 device would make match indices device-dependent, and the whole
// reconstruction with them.
groupshared float g_norm_partial[WG];
[shader("compute")]
[numthreads(WG, 1, 1)]
void l2_normalize_rows(uint3 gid : SV_GroupID, uint3 tid : SV_GroupThreadID,
uniform NormalizeParams p) {
uint row = gid.y * p.groups_per_row + gid.x;
uint t = tid.x;
// Out-of-range rows must still reach every barrier below.
float acc = 0.0f;
if (row < p.rows)
for (uint c = t; c < p.cols; c += uint(WG)) {
float v = load_elem(p.x, row * p.cols + c, kInputF16 != 0);
acc += v * v;
}
g_norm_partial[t] = acc;
GroupMemoryBarrierWithGroupSync();
for (uint s = uint(WG) / 2u; s > 0u; s >>= 1) {
if (t < s) g_norm_partial[t] += g_norm_partial[t + s];
GroupMemoryBarrierWithGroupSync();
}
if (row >= p.rows) return;
float inv = 1.0f / max(sqrt(g_norm_partial[0]), p.eps);
for (uint c = t; c < p.cols; c += uint(WG)) {
float v = load_elem(p.x, row * p.cols + c, kInputF16 != 0);
p.out[row * p.cols + c] = v * inv;
}
}
src/nn/shaders/resample.slang
+71 -7
View File
@@ -4,12 +4,20 @@
// src = (dst + 0.5) * scale - 0.5, clamped, with the low index taken by floor
// and the high index clamped to H-1. Mask upsampling runs through it on every
// frame of a video, and an off-by-half here shifts every boundary by a pixel.
//
// `align_corners=True` is the OTHER mapping -- src = dst * (Hi-1)/(Ho-1) -- and
// is not interchangeable with it: ALIKED's encoder upsamples every block that
// way, and its keypoints are normalized on the same convention, so mixing the
// two moves features by up to half a pixel at every pyramid level. It is a
// specialization constant so each caller states which it means and neither
// pays a branch.
#include "_common.slang"
// Spec ids 0..1.
[SpecializationConstant] const int kInputF16 = 0;
[SpecializationConstant] const int kAct = 0;
// Spec ids 0..2.
[SpecializationConstant] const int kInputF16 = 0;
[SpecializationConstant] const int kAct = 0;
[SpecializationConstant] const int kAlignCorners = 0;
static const int WG = 256;
@@ -24,6 +32,17 @@ struct ResizeParams {
uint groups_per_row;
};
// Source coordinate for output index `o` along one axis.
float src_coord(uint o, uint n_out, uint n_in) {
if (kAlignCorners != 0) {
// A single output sample maps to 0, matching torch (which divides by
// max(n_out - 1, 1) and would otherwise produce a NaN).
if (n_out <= 1u) return 0.0f;
return float(o) * (float(n_in) - 1.0f) / (float(n_out) - 1.0f);
}
return max((float(o) + 0.5f) * (float(n_in) / float(n_out)) - 0.5f, 0.0f);
}
[shader("compute")]
[numthreads(WG, 1, 1)]
void resize_bilinear(uint3 gid : SV_GroupID, uint3 tid : SV_GroupThreadID,
@@ -35,10 +54,8 @@ void resize_bilinear(uint3 gid : SV_GroupID, uint3 tid : SV_GroupThreadID,
uint xo = rest % p.Wo;
uint yo = rest / p.Wo;
float sy = (float(yo) + 0.5f) * (float(p.Hi) / float(p.Ho)) - 0.5f;
float sx = (float(xo) + 0.5f) * (float(p.Wi) / float(p.Wo)) - 0.5f;
sy = max(sy, 0.0f);
sx = max(sx, 0.0f);
float sy = src_coord(yo, p.Ho, p.Hi);
float sx = src_coord(xo, p.Wo, p.Wi);
p.out[i] = sample_bilinear(p.x, kInputF16 != 0, int(p.Hi), int(p.Wi), int(p.C),
int(c), sy, sx);
}
@@ -84,6 +101,53 @@ void maxpool2x2(uint3 gid : SV_GroupID, uint3 tid : SV_GroupThreadID,
p.out[i] = best;
}
// nn.AvgPool2d(kernel, stride), no padding. ALIKED's encoder feeds blocks 3
// and 4 from 2x and 4x pooled copies of the input.
//
// torch with ceil_mode=False drops any window that would run off the edge, so
// every window here is complete and the divisor is exactly kh*kw -- which is
// also why `count_include_pad` never comes into it. The bounds test stays
// because an out-of-range output index must not read, not because a partial
// window is expected.
struct PoolParams {
float* out;
uint* x;
uint Ho;
uint Wo;
uint Hi;
uint Wi;
uint C;
uint kh;
uint kw;
uint stride_y;
uint stride_x;
uint groups_per_row;
};
[shader("compute")]
[numthreads(WG, 1, 1)]
void avgpool(uint3 gid : SV_GroupID, uint3 tid : SV_GroupThreadID,
uniform PoolParams p) {
uint i = fold_index(gid, tid, p.groups_per_row, uint(WG));
if (i >= p.Ho * p.Wo * p.C) return;
uint c = i % p.C;
uint rest = i / p.C;
uint xo = rest % p.Wo;
uint yo = rest / p.Wo;
float sum = 0.0f;
for (uint ky = 0; ky < p.kh; ++ky) {
uint y = yo * p.stride_y + ky;
if (y >= p.Hi) continue;
for (uint kx = 0; kx < p.kw; ++kx) {
uint x = xo * p.stride_x + kx;
if (x >= p.Wi) continue;
sum += load_elem(p.x, (y * p.Wi + x) * p.C + c, kInputF16 != 0);
}
}
p.out[i] = sum / float(p.kh * p.kw);
}
// ---------------------------------------------------------------------------
// Mask export: resize + threshold + pack to bytes, in one pass
// ---------------------------------------------------------------------------
src/nn/tests/nn_ops_test.cpp
+199
View File
@@ -860,6 +860,204 @@ void test_spatial(vk::Arena& arena) {
}
}
// ================
// The learned-frontend ops (src/aliked/)
// ================
// torchvision's bilinear_interpolate, corner-by-corner: zero outside, and a
// corner that falls outside contributes nothing even when the sample point is
// inside. This is the reference the deformable convolution has to match, and
// the reason it is spelled out rather than reusing the clamped sampler above.
float sample_zero(const std::vector<float>& src, int H, int W, int C, int c, float y,
float x) {
if (y <= -1.0f || (float)H <= y || x <= -1.0f || (float)W <= x) return 0.0f;
const int y0 = (int)std::floor(y), x0 = (int)std::floor(x);
const int y1 = y0 + 1, x1 = x0 + 1;
const float ly = y - y0, lx = x - x0, hy = 1 - ly, hx = 1 - lx;
auto at = [&](int yy, int xx) -> float {
if (yy < 0 || yy > H - 1 || xx < 0 || xx > W - 1) return 0.0f;
return src[((size_t)yy * W + xx) * C + c];
};
return hy * (hx * at(y0, x0) + lx * at(y0, x1)) +
ly * (hx * at(y1, x0) + lx * at(y1, x1));
}
void test_learned_frontend(vk::Arena& arena) {
{ // SELU -- ALIKED's gate, in every block and both heads.
vk::ArenaScope scope(arena);
const int N = 4096;
auto x = randn(N, 2.0f);
Tensor to = arena_tensor(arena, DType::F32, N);
unary(to, upload_f32(arena, x, N), Act::Selu);
const float a = 1.6732632423543772848170429916717f;
const float s = 1.0507009873554804934193349852946f;
std::vector<float> want(x.size());
for (size_t i = 0; i < x.size(); ++i)
want[i] = s * (x[i] > 0 ? x[i] : a * (std::exp(x[i]) - 1.0f));
check("selu", readback(to), want, 1e-5f);
}
{ // AvgPool2d(2,2) and (4,4) -- the encoder feeds blocks 3 and 4 from
// pooled copies of the input.
vk::ArenaScope scope(arena);
const int Hi = 16, Wi = 12, C = 5;
auto x = randn((size_t)Hi * Wi * C);
Tensor tx = upload_f32(arena, x, Hi, Wi, C);
for (int k : {2, 4}) {
const int Ho = Hi / k, Wo = Wi / k;
Tensor to = arena_tensor(arena, DType::F32, Ho, Wo, C);
avgpool(to, tx, k);
std::vector<float> want((size_t)Ho * Wo * C, 0.0f);
for (int y = 0; y < Ho; ++y)
for (int xx = 0; xx < Wo; ++xx)
for (int c = 0; c < C; ++c) {
double s = 0;
for (int dy = 0; dy < k; ++dy)
for (int dx = 0; dx < k; ++dx)
s += x[((size_t)(y * k + dy) * Wi + (xx * k + dx)) * C + c];
want[((size_t)y * Wo + xx) * C + c] = (float)(s / (k * k));
}
check(k == 2 ? "avgpool 2x2" : "avgpool 4x4", readback(to), want, 1e-5f);
}
}
{ // resize_bilinear(align_corners=True): the OTHER coordinate mapping.
vk::ArenaScope scope(arena);
const int Hi = 5, Wi = 7, C = 3, Ho = 17, Wo = 11;
auto x = randn((size_t)Hi * Wi * C);
Tensor to = arena_tensor(arena, DType::F32, Ho, Wo, C);
resize_bilinear(to, upload_f32(arena, x, Hi, Wi, C), /*align_corners=*/true);
std::vector<float> want((size_t)Ho * Wo * C);
for (int y = 0; y < Ho; ++y)
for (int xx = 0; xx < Wo; ++xx) {
const float sy = (float)y * (Hi - 1.0f) / (Ho - 1.0f);
const float sx = (float)xx * (Wi - 1.0f) / (Wo - 1.0f);
int y0 = (int)std::floor(sy), x0 = (int)std::floor(sx);
const float fy = sy - y0, fx = sx - x0;
const int y1 = std::min(y0 + 1, Hi - 1), x1 = std::min(x0 + 1, Wi - 1);
y0 = std::min(y0, Hi - 1);
x0 = std::min(x0, Wi - 1);
for (int c = 0; c < C; ++c)
want[((size_t)y * Wo + xx) * C + c] =
(1 - fy) * ((1 - fx) * x[((size_t)y0 * Wi + x0) * C + c] +
fx * x[((size_t)y0 * Wi + x1) * C + c]) +
fy * ((1 - fx) * x[((size_t)y1 * Wi + x0) * C + c] +
fx * x[((size_t)y1 * Wi + x1) * C + c]);
}
check("resize_bilinear align_corners", readback(to), want, 1e-4f);
}
{ // deform_conv2d vs torchvision's definition. The offsets are large on
// purpose -- +-3 px on a 9x11 map -- so a good fraction of the taps
// land outside and the zero-fill rule is actually exercised.
vk::ArenaScope scope(arena);
const int Hi = 9, Wi = 11, Ci = 6, Co = 7, k = 3, pad = 1;
const int Ho = Hi, Wo = Wi;
const float max_offset = 2.5f;
auto x = randn((size_t)Hi * Wi * Ci);
auto w = randn((size_t)Co * Ci * k * k, 0.3f);
auto bias = randn(Co);
auto off = randn((size_t)Ho * Wo * 2 * k * k, 3.0f);
Tensor tx = upload_f32(arena, x, Hi, Wi, Ci);
Tensor tw = upload_f32(arena, w, Co, Ci * k * k);
Tensor tb = upload_f32(arena, bias, Co);
Tensor toff = upload_f32(arena, off, Ho, Wo, 2 * k * k);
Tensor to = arena_tensor(arena, DType::F32, Ho, Wo, Co);
ConvOpts o;
o.pad_y = o.pad_x = pad;
o.bias = tb;
deform_conv2d(arena, to, tx, toff, tw, k, k, max_offset, o);
std::vector<float> want((size_t)Ho * Wo * Co);
for (int yo = 0; yo < Ho; ++yo)
for (int xo = 0; xo < Wo; ++xo)
for (int co = 0; co < Co; ++co) {
double s = bias[(size_t)co];
for (int ci = 0; ci < Ci; ++ci)
for (int ky = 0; ky < k; ++ky)
for (int kx = 0; kx < k; ++kx) {
const int tap = ky * k + kx;
const size_t ob =
((size_t)yo * Wo + xo) * 2 * k * k + 2 * tap;
const float dy = std::fmax(
-max_offset, std::fmin(max_offset, off[ob]));
const float dx = std::fmax(
-max_offset, std::fmin(max_offset, off[ob + 1]));
const float sy = (float)(yo + ky - pad) + dy;
const float sx = (float)(xo + kx - pad) + dx;
s += (double)w[((size_t)co * Ci + ci) * k * k + tap] *
sample_zero(x, Hi, Wi, Ci, ci, sy, sx);
}
want[((size_t)yo * Wo + xo) * Co + co] = (float)s;
}
check("deform_conv2d", readback(to), want, 2e-4f);
}
{ // patch_gather: a k x k patch per centre, in conv-weight column order,
// with centres deliberately on and past the border.
vk::ArenaScope scope(arena);
const int Hi = 7, Wi = 9, C = 4, k = 3, N = 5;
auto x = randn((size_t)Hi * Wi * C);
const int32_t centers[N * 2] = {3, 3, 0, 0, 8, 6, 4, 0, 0, 6};
Tensor tc = arena_tensor(arena, DType::I32, N, 2);
vk::Stream::get().upload(tc.ptr, centers, sizeof centers);
Tensor to = arena_tensor(arena, DType::F32, N, C * k * k);
patch_gather(to, upload_f32(arena, x, Hi, Wi, C), tc, k);
std::vector<float> want((size_t)N * C * k * k, 0.0f);
for (int n = 0; n < N; ++n)
for (int c = 0; c < C; ++c)
for (int ky = 0; ky < k; ++ky)
for (int kx = 0; kx < k; ++kx) {
const int y = centers[2 * n + 1] - k / 2 + ky;
const int xx = centers[2 * n] - k / 2 + kx;
if (y < 0 || y >= Hi || xx < 0 || xx >= Wi) continue;
want[(size_t)n * C * k * k + (size_t)c * k * k + ky * k + kx] =
x[((size_t)y * Wi + xx) * C + c];
}
check("patch_gather", readback(to), want, 1e-6f);
}
{ // grid_sample_points, both conventions, with points outside the map so
// the zero padding is covered too.
vk::ArenaScope scope(arena);
const int H = 8, W = 6, C = 3, N = 64;
auto x = randn((size_t)H * W * C);
auto pos = randn((size_t)N * 2, 0.8f);
Tensor tx = upload_f32(arena, x, H, W, C);
Tensor tp = upload_f32(arena, pos, N, 2);
for (bool ac : {true, false}) {
Tensor to = arena_tensor(arena, DType::F32, N, C);
grid_sample_points(to, tx, tp, ac);
std::vector<float> want((size_t)N * C);
for (int n = 0; n < N; ++n) {
const float gx = pos[(size_t)n * 2], gy = pos[(size_t)n * 2 + 1];
const float sx = ac ? (gx + 1) * 0.5f * (W - 1) : ((gx + 1) * W - 1) * 0.5f;
const float sy = ac ? (gy + 1) * 0.5f * (H - 1) : ((gy + 1) * H - 1) * 0.5f;
for (int c = 0; c < C; ++c)
want[(size_t)n * C + c] = sample_zero(x, H, W, C, c, sy, sx);
}
check(ac ? "grid_sample_points align" : "grid_sample_points noalign",
readback(to), want, 1e-5f);
}
}
{ // l2_normalize_rows, including a row of exact zeros (a keypoint whose
// descriptor cancelled) -- eps must keep that finite, not NaN.
vk::ArenaScope scope(arena);
const int rows = 300, cols = 128;
auto x = randn((size_t)rows * cols);
for (int c = 0; c < cols; ++c) x[(size_t)7 * cols + c] = 0.0f;
Tensor to = arena_tensor(arena, DType::F32, rows, cols);
l2_normalize_rows(to, upload_f32(arena, x, rows, cols));
std::vector<float> want((size_t)rows * cols);
for (int r = 0; r < rows; ++r) {
double sq = 0;
for (int c = 0; c < cols; ++c) sq += (double)x[(size_t)r * cols + c] *
x[(size_t)r * cols + c];
const float inv = 1.0f / std::fmax((float)std::sqrt(sq), 1e-12f);
for (int c = 0; c < cols; ++c)
want[(size_t)r * cols + c] = x[(size_t)r * cols + c] * inv;
}
check("l2_normalize_rows", readback(to), want, 1e-5f);
}
}
// ================
// Benchmark (test_ops --bench)
// ================
@@ -1002,6 +1200,7 @@ int main(int argc, char** argv) {
std::printf("RoPE\n"); test_rope(arena);
std::printf("Convolution\n"); test_conv(arena);
std::printf("Spatial / gather\n"); test_spatial(arena);
std::printf("Learned frontend\n"); test_learned_frontend(arena);
std::printf("\n%d checks, %d failures\n", g_checks, g_failures);
} catch (const std::exception& e) {
src/sfm/SfmConfig.cpp
+83 -7
View File
@@ -157,8 +157,15 @@ constexpr size_t kInlineChoices = 26;
std::string metavarFor(bool, const char*, const char*) { return ""; }
std::string metavarFor(const std::string&, const char* name, const char* choices) {
if (choices && *choices && std::string(choices).size() <= kInlineChoices)
return "{" + std::string(choices) + "}";
if (choices && *choices) {
// A field WITH choices never takes a path, whatever it is called --
// `--features` is a directory on `map` and a frontend name on
// `extract`, and only the latter has choices. Long lists do not fit
// the flag column and are printed in the help text instead.
return std::string(choices).size() <= kInlineChoices
? "{" + std::string(choices) + "}"
: "VALUE";
}
// No metavar column in the table: the flag name says what it takes, and
// "DIR" reads better than "VALUE" on the handful that take a path.
std::string n = name;
@@ -274,25 +281,75 @@ std::string applyPresets(SfmConfig& cfg, const std::set<std::string>& seen,
// of the extractor's 3200 px default. Pair-selection breadth follows it too
// (D42): `prefilter-neighbors` is the one selection parameter that trades
// match time against how much of the view graph verification even sees.
// The learned frontend runs on its own resolution ladder, roughly two
// thirds of SIFT's -- COLMAP does the same (1600 against 3200 at the
// default), and for the same reason: ALIKED aggregates 128 channels at
// FULL resolution, so working size is what its memory is spent on. Its
// feature counts are lower too, which is the detector's design rather
// than a budget: it emits fewer, better-localized points.
const bool learned = isAlikedType(cfg.features);
if (cfg.quality == "low") {
presetSet(seen, moved, "max-image-size", cfg.max_image_size, 1000);
presetSet(seen, moved, "max-image-size", cfg.max_image_size, learned ? 800 : 1000);
presetSet(seen, moved, "max-features", cfg.sift.max_num_features, 2048);
presetSet(seen, moved, "aliked-max-features", cfg.aliked.max_num_features, 1024);
presetSet(seen, moved, "prefilter-neighbors", cfg.prefilter.num_neighbors, 16);
} else if (cfg.quality == "medium") {
presetSet(seen, moved, "max-image-size", cfg.max_image_size, 1600);
presetSet(seen, moved, "max-image-size", cfg.max_image_size, learned ? 1200 : 1600);
presetSet(seen, moved, "max-features", cfg.sift.max_num_features, 4096);
presetSet(seen, moved, "aliked-max-features", cfg.aliked.max_num_features, 2048);
presetSet(seen, moved, "prefilter-neighbors", cfg.prefilter.num_neighbors, 24);
} else if (cfg.quality == "high") {
presetSet(seen, moved, "max-image-size", cfg.max_image_size, 2400);
presetSet(seen, moved, "max-image-size", cfg.max_image_size, learned ? 1600 : 2400);
presetSet(seen, moved, "max-features", cfg.sift.max_num_features, 8192);
presetSet(seen, moved, "aliked-max-features", cfg.aliked.max_num_features, 4096);
} else if (cfg.quality == "extreme") {
presetSet(seen, moved, "max-image-size", cfg.max_image_size, 3200);
presetSet(seen, moved, "max-image-size", cfg.max_image_size, learned ? 2400 : 3200);
presetSet(seen, moved, "max-features", cfg.sift.max_num_features, 16384);
presetSet(seen, moved, "aliked-max-features", cfg.aliked.max_num_features, 8192);
presetSet(seen, moved, "prefilter-neighbors", cfg.prefilter.num_neighbors, 48);
} else {
return "unknown --quality '" + cfg.quality + "' (low, medium, high or extreme)";
}
// A learned descriptor needs a looser ratio than SIFT's 0.8, because its
// second-best distance sits much closer to its best: measured on a
// 20-image capture, the median mutual-nearest match has a distance ratio
// of 0.826, i.e. just the wrong side of SIFT's threshold. The whole point
// of the ratio test -- that a true match stands out from the runner-up --
// is weaker for descriptors trained to be smooth, which is also why
// LightGlue exists.
//
// 0.92 is where the measurement put it, on 190 exhaustive pairs:
//
// ratio 0.80 0.85 0.90 0.92 0.95
// pairs 65/190 79 111 172 190
// inliers 6331 8426 11487 13633 16249
// inlier% 90 79 54 44 29
//
// Past 0.92 the pair count is bought with junk. For reference, GPU SIFT on
// the same images at 4x the feature budget managed 68/190 and 7703
// inliers, so this is not a concession -- it is where the learned frontend
// wins.
//
// COLMAP's own ALIKED defaults (ratio off, min_cossim 0.85) were tried
// first and are NOT used: on this data an absolute 0.85 cosine rejects
// ~70% of mutual-nearest matches (their median cosine is 0.726) and left
// 24/190 pairs. --min-similarity still exists for anyone who wants that
// shape of test.
if (learned && !isLearnedMatcher(cfg.matcher))
presetSet(seen, moved, "ratio", cfg.match.max_ratio, 0.92f);
// LightGlue decides on its own assignment, so the ratio test and the
// cross-check are not merely unnecessary -- they are a second filter on a
// quantity it does not produce. Its own confidence is the only threshold.
if (isLearnedMatcher(cfg.matcher)) {
presetSet(seen, moved, "ratio", cfg.match.max_ratio, 1.0f);
presetSet(seen, moved, "min-similarity", cfg.match.min_similarity, 0.0f);
// Exhaustive matching with it is hours on a capture where pair
// selection is minutes, and the shortlist is what it was built for.
presetSet(seen, moved, "pairs", cfg.pairs, std::string("prefilter"));
}
if (cfg.data_type == "individual") {
// Nothing: the defaults are written for a set of individual photos.
} else if (cfg.data_type == "video") {
@@ -332,10 +389,29 @@ std::string SfmConfig::finalize(uint32_t cmd) {
twoview.ransac.max_error = max_error;
mapper.max_reproj_error = max_error;
if (features != "sift" && !isAlikedType(features))
return "unknown --features '" + features +
"' (sift, aliked-n16rot or aliked-n32)";
if (matcher != "bruteforce" && !isLearnedMatcher(matcher))
return "unknown --matcher '" + matcher + "' (bruteforce or lightglue)";
// LightGlue is trained on one frontend's descriptors, and matching SIFT
// with it would run and return nonsense. Only `auto` can check that here:
// `match` reads features off disk and has no --features to compare
// against, so its guard is on the descriptors themselves, in
// LearnedMatcher.cpp, which is the more honest place for it anyway.
if ((cmd & (CMD_AUTO | CMD_EXTRACT)) && isLearnedMatcher(matcher) &&
!isAlikedType(features))
return "--matcher " + matcher + " needs learned descriptors; add "
"--features aliked-n16rot";
lightglue.device = device;
if (max_image_size <= 0) max_image_size = defaultMaxImageSize(features);
sift.device = match.device = prefilter.device = mapper.device = device;
aliked.device = device;
mapper.threads = threads;
const bool v = !quiet;
sift.verbose = mapper.verbose = manager.verbose = merge.verbose = v;
sift.verbose = mapper.verbose = manager.verbose = merge.verbose = aliked.verbose = v;
lightglue.verbose = v;
// Scoring problems are ~1/32 the size of full matching, so the selection
// pass batches at least as many pairs per submit as the matcher does.
src/sfm/SfmConfig.h
+47 -6
View File
@@ -32,6 +32,8 @@
#include "sfm/core/CameraSetup.h"
#include "sfm/feature/Matcher.h"
#include "sfm/feature/Extractor.h"
#include "sfm/feature/LearnedMatcher.h"
#include "sfm/feature/PairSelection.h"
#include "sfm/feature/Sift.h"
#include "sfm/geometry/TwoView.h"
@@ -76,7 +78,12 @@ struct SfmConfig {
// verifier's inlier radius and the mapper's reprojection cap are the same
// quantity in the same frame. finalize() writes both.
double max_error = 3.0;
int max_image_size = 3200;
// 0 means "whatever the selected frontend wants" -- 3200 for SIFT, 1600
// for a learned one, mirroring COLMAP's EffMaxImageSize(). Resolved in
// finalize(), so a command that applies no presets (extract, match) still
// gets the right one instead of running ALIKED at SIFT's resolution and
// spending four times the VRAM on it.
int max_image_size = 0;
std::string mask_dir;
// Camera setup. The string forms are what the table and the GUI see; the
@@ -111,8 +118,18 @@ struct SfmConfig {
int device = -1;
bool quiet = false;
// Which frontend runs. "sift" is the GPU SIFT that has always been here;
// the aliked-* values are the learned one (src/aliked/), which needs the
// inference layer compiled in. The pair is deliberately two flags and not
// one: ALIKED descriptors can be matched brute-force, and LightGlue is a
// matcher for them rather than a different extractor.
std::string features = "sift";
std::string matcher = "bruteforce";
// ---- the stage option structs, unchanged ----
SiftOptions sift;
AlikedOptions aliked;
LightGlueOptions lightglue;
MatchOptions match;
PairSelectionOptions prefilter;
TwoViewOptions twoview;
@@ -183,9 +200,10 @@ struct SfmConfig {
"Inlier radius for verification and mapping, in pixels of the image SIFT ran on rather " \
"than of the source file (D47)") \
F(max_image_size, "max-image-size", CMD_AUTO | CMD_EXTRACT, Tier::Advanced, "pipeline", \
64, 20000, "", \
"Longest edge SIFT runs on; larger images are downscaled first, and keypoints are still " \
"reported in the source image's pixels") \
0, 20000, "", \
"Longest edge the extractor runs on; larger images are downscaled first, and keypoints are " \
"still reported in the source image's pixels. 0 picks the frontend's own default (3200 for " \
"sift, 1600 for aliked)") \
F(mask_dir, "masks", CMD_AUTO | CMD_EXTRACT, Tier::Basic, "pipeline", 0, 0, "", \
"Directory of masks; keypoints on zero (black) pixels are dropped. auto defaults it to " \
"`masks` beside the image directory") \
@@ -211,9 +229,20 @@ struct SfmConfig {
"camera", 0.001, 1.0, "", \
"Relative tolerance clustering EXIF focals into one group; must exceed EXIF's 1 mm " \
"quantization and stay under a real zoom step") \
/* ---- features (SIFT) ---- */ \
/* ---- features ---- */ \
F(features, "features", CMD_AUTO | CMD_EXTRACT, Tier::Basic, "features", 0, 0, \
"sift|aliked-n16rot|aliked-n32", \
"Which detector and descriptor; the aliked ones are learned and fetch a checkpoint on " \
"first use") \
F(sift.max_num_features, "max-features", CMD_AUTO | CMD_EXTRACT, Tier::Advanced, "features", \
128, 1000000, "", "Keypoints kept per image, the largest scales first") \
128, 1000000, "", "Keypoints kept per image with --features sift, the largest scales first") \
F(aliked.max_num_features, "aliked-max-features", CMD_AUTO | CMD_EXTRACT, Tier::Advanced, \
"features", 128, 1000000, "", \
"Keypoints kept per image with a learned frontend, the highest scores first") \
F(aliked.min_score, "aliked-min-score", CMD_AUTO | CMD_EXTRACT, Tier::Advanced, "features", \
0, 1, "", "Detection score a learned keypoint must reach") \
F(aliked.model, "aliked-model", CMD_AUTO | CMD_EXTRACT, Tier::Advanced, "features", 0, 0, "", \
"Path to an ALIKED .onnx checkpoint, overriding the one --features names") \
F(sift.num_octaves, "octaves", CMD_AUTO | CMD_EXTRACT, Tier::Advanced, "features", 1, 8, "", \
"Scale-space octaves") \
F(sift.peak_threshold, "peak-threshold", CMD_AUTO | CMD_EXTRACT, Tier::Advanced, "features", \
@@ -227,8 +256,20 @@ struct SfmConfig {
F(sift.spv_path, "spv-path", CMD_EXTRACT, Tier::Advanced, "features", 0, 0, "", \
"Load the SIFT kernels from this SPIR-V file instead of the embedded blob") \
/* ---- matching ---- */ \
F(matcher, "matcher", CMD_AUTO | CMD_MATCH, Tier::Basic, "matching", 0, 0, \
"bruteforce|lightglue", \
"How descriptors are matched. lightglue is a learned matcher for --features aliked-*; it is " \
"an order of magnitude slower per pair, so it only makes sense behind pair selection") \
F(lightglue.min_score, "lightglue-min-score", CMD_AUTO | CMD_MATCH, Tier::Advanced, \
"matching", 0, 1, "", "Assignment confidence a LightGlue match must reach") \
F(lightglue.model, "lightglue-model", CMD_AUTO | CMD_MATCH, Tier::Advanced, "matching", \
0, 0, "", "Path to a LightGlue .onnx checkpoint, overriding the fetched one") \
F(match.max_ratio, "ratio", CMD_AUTO | CMD_MATCH, Tier::Advanced, "matching", 0, 1, "", \
"Lowe ratio: a match is kept when the best distance is below this times the second best") \
F(match.min_similarity, "min-similarity", CMD_AUTO | CMD_MATCH, Tier::Advanced, "matching", \
0, 1, "", \
"Cosine similarity a match must reach, for float descriptors; 0 disables it. Learned " \
"descriptors are filtered on this rather than on the ratio") \
F(match.cross_check, "cross-check", CMD_AUTO | CMD_MATCH, Tier::Advanced, "matching", 0, 0, \
"", "Keep only mutual nearest neighbours") \
F(match.max_num_matches, "max-matches", CMD_AUTO | CMD_MATCH, Tier::Advanced, "matching", \
src/sfm/core/Features.h
+37 -2
View File
@@ -25,7 +25,11 @@ struct Keypoint {
float x = 0, y = 0; // subpixel location
float scale = 0; // sigma in original-image pixels
float orientation = 0; // radians, CCW from +x
float response = 0; // |DoG| at the refined extremum (not persisted)
// SIFT: |DoG| at the refined extremum. A learned detector: its detection
// score. Persisted from v5 -- ALIKED has neither scale nor orientation, so
// this is the only ranking signal its keypoints carry, and anything that
// used to rank by scale has to fall back to it (Pairing / PairSelection).
float response = 0;
};
enum class DType : uint32_t { U8 = 0, F32 = 1 };
@@ -61,6 +65,20 @@ struct FeatureSet {
uint32_t count() const { return (uint32_t)keypoints.size(); }
bool hasColors() const { return colors.size() == (size_t)count() * 3; }
// Whether any keypoint carries a detection score worth persisting. SIFT
// leaves response at 0; a learned detector fills it in.
bool hasScores() const {
for (const Keypoint& k : keypoints)
if (k.response != 0) return true;
return false;
}
// What a subset selection should rank by. Scale for SIFT (D16: the largest
// scales are the most repeatable), the detection score for a detector that
// has no scale. Never both -- an extractor fills in one of them.
float rank(uint32_t i) const {
const Keypoint& k = keypoints[i];
return k.scale > 0 ? k.scale : k.response;
}
// How many source pixels one extraction pixel is worth: >= 1, and exactly 1
// when nothing was downscaled. The two axes agree up to the rounding in
// the size clamp, so their mean is the isotropic answer.
@@ -116,11 +134,15 @@ inline void scaleKeypoints(FeatureSet& fs, int w, int h) {
// v4 appends i32 extract_width, extract_height after that. Older files read
// back with 0, i.e. pixelScale() == 1, i.e. thresholds in source pixels --
// which is what those files were produced under.
//
// v5 appends u8 has_scores, then (if 1) count f32 detection scores. A v1-v4
// file reads back with every response 0, which is what those files carried:
// SIFT never persisted it and nothing read it.
inline void writeFeatures(const std::string& path, const FeatureSet& fs) {
std::ofstream f(path, std::ios::binary);
if (!f) throw std::runtime_error("cannot write " + path);
uint32_t version = 4, count = fs.count(), dtype = (uint32_t)fs.dtype;
uint32_t version = 5, count = fs.count(), dtype = (uint32_t)fs.dtype;
f.write("VKFT", 4);
f.write((const char*)&version, 4);
f.write((const char*)&fs.width, 4);
@@ -142,6 +164,10 @@ inline void writeFeatures(const std::string& path, const FeatureSet& fs) {
f.write(fs.exif_camera.data(), (std::streamsize)cam_len);
f.write((const char*)&fs.extract_width, 4);
f.write((const char*)&fs.extract_height, 4);
uint8_t has_scores = fs.hasScores() ? 1 : 0;
f.write((const char*)&has_scores, 1);
if (has_scores)
for (const Keypoint& k : fs.keypoints) f.write((const char*)&k.response, 4);
}
// `with_descriptors == false` seeks past the descriptor block instead of
@@ -217,6 +243,15 @@ inline FeatureSet readFeatures(const std::string& path, bool with_descriptors =
fs.extract_height = eh;
}
}
if (version >= 5) {
uint8_t has_scores = 0;
f.read((char*)&has_scores, 1);
if (f.gcount() == 1 && has_scores) {
std::vector<float> raw(count);
f.read((char*)raw.data(), (std::streamsize)(raw.size() * sizeof(float)));
for (uint32_t i = 0; i < count; i++) fs.keypoints[i].response = raw[i];
}
}
return fs;
}
src/sfm/feature/Extractor.cpp
+128
View File
@@ -0,0 +1,128 @@
#include "sfm/feature/Extractor.h"
#include <algorithm>
#include <cstring>
#include <stdexcept>
#include <vector>
#if SSPLAT_HAVE_ALIKED
#include "aliked/Aliked.h"
#endif
namespace sfm {
namespace {
class SiftFrontend : public IFeatureExtractor {
public:
explicit SiftFrontend(const SiftOptions& opt) : ext_(opt) {}
FeatureSet extract(const GrayImage& img) override { return ext_.extract(img); }
const char* name() const override { return "sift"; }
private:
SiftExtractor ext_;
};
#if SSPLAT_HAVE_ALIKED
// ALIKED behind the same contract.
//
// The conversions in both directions are the whole of it, and each is a
// convention this repository already fixed elsewhere:
//
// * input -- the loader's optional RGB companion buffer, which `extract`
// already asks for so it can colour the point cloud. ALIKED normalizes by
// 1/255 and nothing else, so the bytes go straight through.
// * output -- keypoints have no scale and no orientation, because the
// detector has neither. They carry a detection score instead, which
// `FeatureSet::rank` and features.bin v5 exist for. Leaving `scale` at 0
// is deliberate: a fabricated scale would be silently wrong wherever one
// is used as a size.
class AlikedFrontend : public IFeatureExtractor {
public:
explicit AlikedFrontend(const AlikedOptions& opt) : opt_(opt) {
ext_.load(opt.model);
aopts_.max_num_features = opt.max_num_features;
aopts_.min_score = (float)opt.min_score;
}
const char* name() const override { return "aliked"; }
bool wantsColor() const override { return true; }
FeatureSet extract(const GrayImage& img) override {
FeatureSet fs;
fs.width = img.width;
fs.height = img.height;
fs.dim = (uint32_t)ext_.descriptorDim();
fs.dtype = DType::F32;
if (!img.hasColor())
throw std::runtime_error(
"ALIKED needs a colour image; the loader decoded luma only");
const aliked::Features f =
ext_.extract(img.rgb.data(), img.width, img.height, aopts_);
// Canonical order by position, exactly as GPU SIFT emits (D16). The
// extractor's own order comes from a partial_sort by score, and a
// score-sorted index would bias every downstream tie-break -- the
// matcher's cap, the mapper taking the first 3D point a feature
// corresponds to -- toward high-scoring features.
std::vector<uint32_t> idx(f.keypoints.size());
for (uint32_t i = 0; i < idx.size(); i++) idx[i] = i;
std::sort(idx.begin(), idx.end(), [&](uint32_t a, uint32_t b) {
const aliked::Keypoint& p = f.keypoints[a];
const aliked::Keypoint& q = f.keypoints[b];
if (p.x != q.x) return p.x < q.x;
if (p.y != q.y) return p.y < q.y;
return p.score > q.score;
});
const size_t n = idx.size();
fs.keypoints.resize(n);
fs.descriptors.resize(n * fs.dim * sizeof(float));
float* dst = reinterpret_cast<float*>(fs.descriptors.data());
for (size_t i = 0; i < n; i++) {
const aliked::Keypoint& k = f.keypoints[idx[i]];
fs.keypoints[i] = {k.x, k.y, /*scale=*/0.0f, /*orientation=*/0.0f, k.score};
std::memcpy(dst + i * fs.dim, &f.descriptors[(size_t)idx[i] * fs.dim],
fs.dim * sizeof(float));
}
if (opt_.verbose) fprintf(stderr, "[aliked] %zu features\n", n);
return fs;
}
private:
AlikedOptions opt_;
aliked::Extractor ext_;
aliked::ExtractOptions aopts_;
};
#endif // SSPLAT_HAVE_ALIKED
} // namespace
bool isAlikedType(const std::string& type) { return type.rfind("aliked", 0) == 0; }
int defaultMaxImageSize(const std::string& type) { return isAlikedType(type) ? 1600 : 3200; }
std::unique_ptr<IFeatureExtractor> createFeatureExtractor(const std::string& type,
const SiftOptions& sift,
const AlikedOptions& aliked) {
if (type == "sift") return std::make_unique<SiftFrontend>(sift);
if (isAlikedType(type)) {
#if SSPLAT_HAVE_ALIKED
AlikedOptions opt = aliked;
// --features names the checkpoint, so an explicit --aliked-model is
// only needed to point at a file on disk.
if (opt.model.empty() || isAlikedType(opt.model)) opt.model = type;
return std::make_unique<AlikedFrontend>(opt);
#else
throw std::runtime_error(
"this build has no learned frontend: '" + type +
"' needs the inference layer, which is SSPLAT_BUILD_SAM=ON");
#endif
}
throw std::runtime_error("unknown feature type '" + type +
"' (expected sift, aliked-n16rot or aliked-n32)");
}
} // namespace sfm
src/sfm/feature/Extractor.h
+69
View File
@@ -0,0 +1,69 @@
// The extractor seam: what `extract` calls, and the factory that decides which
// one it gets.
//
// GPU SIFT (sfm/feature/Sift.h) was the only implementation and the extract
// stage constructed it directly. That is no longer true -- ALIKED
// (src/aliked/) implements the same contract -- and the indirection buys two
// things beyond the obvious:
//
// * the SIFT extractor's Vulkan context is not created when it is not the
// one selected, which matters because ALIKED carries the inference layer's
// own device and two live devices on one GPU is something this repository
// sequences deliberately (AGENTS.md);
// * each extractor states what it needs from the loader (color, working
// resolution) rather than the CLI knowing per-type rules.
//
// Nothing above this seam assumes 128-D uint8, a scale or an orientation.
#pragma once
#include <memory>
#include <string>
#include "sfm/core/Features.h"
#include "sfm/core/Image.h"
#include "sfm/feature/Sift.h"
namespace sfm {
// ALIKED's knobs. Deliberately not aliked::ExtractOptions: src/sfm/ must build
// without the inference layer, so nothing here may include an aliked/ header.
struct AlikedOptions {
// "aliked-n16rot" or "aliked-n32" (fetched and cached on first use), or a
// path to an .onnx file. The two released variants differ only in how many
// sample positions the descriptor head uses.
std::string model = "aliked-n16rot";
int max_num_features = 2048; // COLMAP's AlikedExtractionOptions default
double min_score = 0.2;
bool verbose = true;
int device = -1;
};
struct IFeatureExtractor {
virtual ~IFeatureExtractor() = default;
// One image in, its features out, in the coordinates of `img` (the caller
// scales them back to the source file's, D46).
virtual FeatureSet extract(const GrayImage& img) = 0;
virtual const char* name() const = 0;
// Whether the loader has to decode color. SIFT works on luma; a learned
// detector was trained on RGB and must see it.
virtual bool wantsColor() const { return false; }
};
// Longest edge this extractor should run at when the user did not say, mirroring
// COLMAP's FeatureExtractionOptions::EffMaxImageSize(): 3200 for SIFT, 1600 for
// ALIKED, whose aggregated feature map is 128 channels at full resolution.
int defaultMaxImageSize(const std::string& type);
// Whether `type` names a learned frontend at all (as opposed to "sift").
bool isAlikedType(const std::string& type);
// Throws std::runtime_error naming the type when it is unknown, or when it is
// ALIKED and this binary was built without the inference layer.
std::unique_ptr<IFeatureExtractor> createFeatureExtractor(const std::string& type,
const SiftOptions& sift,
const AlikedOptions& aliked);
} // namespace sfm
src/sfm/feature/LearnedMatcher.cpp
+115
View File
@@ -0,0 +1,115 @@
#include "sfm/feature/LearnedMatcher.h"
#include <algorithm>
#include <cstdio>
#include <stdexcept>
#include <vector>
#if SSPLAT_HAVE_ALIKED
#include "aliked/model/LightGlue.h"
#endif
namespace sfm {
namespace {
#if SSPLAT_HAVE_ALIKED
class LightGlueMatcher : public IFeatureMatcher {
public:
explicit LightGlueMatcher(const MatchOptions& match, const LightGlueOptions& opt)
: match_(match), opt_(opt) {
lg_.load(opt.model);
mopt_.min_score = (float)opt.min_score;
}
const char* name() const override { return "lightglue"; }
std::vector<FeatureMatch> match(const FeatureSet& a, const FeatureSet& b) override {
std::vector<FeatureMatch> out;
if (a.count() == 0 || b.count() == 0) return out;
check(a);
check(b);
std::vector<float> ka, kb;
const std::vector<aliked::Match> m =
lg_.match(view(a, ka), view(b, kb), mopt_);
out.reserve(m.size());
for (const aliked::Match& x : m) {
// FeatureMatch carries a *distance*, and everything downstream that
// reads it sorts ascending (the max_num_matches cap). LightGlue
// reports a confidence in (0, 1], so the distance is its
// complement -- a monotone map, which is all that ordering needs.
out.push_back({x.i, x.j, 1.0f - x.score});
}
if (match_.max_num_matches > 0 && out.size() > match_.max_num_matches) {
std::partial_sort(out.begin(), out.begin() + match_.max_num_matches, out.end(),
[](const FeatureMatch& p, const FeatureMatch& q) {
return p.distance < q.distance;
});
out.resize(match_.max_num_matches);
}
return out;
}
private:
static void check(const FeatureSet& f) {
if (f.dtype != DType::F32)
throw std::runtime_error(
"LightGlue needs float descriptors; these are uint8. It is trained for "
"a specific frontend -- use --features aliked-n16rot or aliked-n32");
}
// The keypoints LightGlue sees must be in the frame its `image_size` names.
// FeatureSet reports them in the SOURCE image's pixels (D46), and
// width/height are that same frame, so the two already agree -- but only
// because scaleKeypoints ran. Passing extract_width here instead would
// scale every keypoint by the downscale factor and quietly ruin the
// positional encoding.
// `xy` is filled with the packed [n, 2] the model wants: Keypoint carries
// five floats, so its x/y are strided, not an array of pairs.
static aliked::MatchInput view(const FeatureSet& f, std::vector<float>& xy) {
xy.resize((size_t)f.count() * 2);
for (uint32_t i = 0; i < f.count(); i++) {
xy[(size_t)i * 2] = f.keypoints[i].x;
xy[(size_t)i * 2 + 1] = f.keypoints[i].y;
}
aliked::MatchInput in;
in.keypoints = xy.data();
in.descriptors = reinterpret_cast<const float*>(f.descriptors.data());
in.n = f.count();
in.width = f.width;
in.height = f.height;
return in;
}
MatchOptions match_;
LightGlueOptions opt_;
aliked::Matcher lg_;
aliked::MatchOptions mopt_;
};
#endif // SSPLAT_HAVE_ALIKED
} // namespace
bool isLearnedMatcher(const std::string& type) { return type == "lightglue"; }
std::unique_ptr<IFeatureMatcher> createFeatureMatcher(const std::string& type,
const MatchOptions& match,
const LightGlueOptions& lightglue) {
if (type == "bruteforce") return std::make_unique<BruteForceMatcher>(match);
if (type == "lightglue") {
#if SSPLAT_HAVE_ALIKED
return std::make_unique<LightGlueMatcher>(match, lightglue);
#else
throw std::runtime_error(
"this build has no learned matcher: --matcher lightglue needs the "
"inference layer, which is SSPLAT_BUILD_SAM=ON");
#endif
}
throw std::runtime_error("unknown matcher '" + type + "' (expected bruteforce or "
"lightglue)");
}
} // namespace sfm
src/sfm/feature/LearnedMatcher.h
+41
View File
@@ -0,0 +1,41 @@
#pragma once
// The matcher seam: LightGlue behind IFeatureMatcher, and the factory that
// picks between it and the GPU brute-force matcher.
//
// sfm/feature/Matcher.h always said a learned matcher would implement that
// interface; this is it. What the interface does NOT carry, and what callers
// therefore have to know, is cost: brute force is ~3.5 ms for a pair, LightGlue
// is tens of milliseconds, because it runs nine transformer layers over both
// images' keypoints. It belongs behind pair selection and nowhere else --
// `--pairs exhaustive` with `--matcher lightglue` on a thousand images is half
// a million pairs and several hours.
//
// Like Extractor.h, nothing here may include an aliked/ header: src/sfm/ builds
// without the inference layer, and the factory says so at run time.
#include <memory>
#include <string>
#include "sfm/feature/Matcher.h"
namespace sfm {
struct LightGlueOptions {
// "aliked-lightglue" (fetched and cached on first use) or a path to an
// .onnx file.
std::string model = "aliked-lightglue";
// COLMAP's LightGlueONNXMatchingOptions default.
double min_score = 0.1;
int device = -1;
bool verbose = true;
};
bool isLearnedMatcher(const std::string& type);
// Throws std::runtime_error naming the type when it is unknown, or when it is
// learned and this binary has no inference layer.
std::unique_ptr<IFeatureMatcher> createFeatureMatcher(const std::string& type,
const MatchOptions& match,
const LightGlueOptions& lightglue);
} // namespace sfm
src/sfm/feature/Matcher.h
+54 -5
View File
@@ -29,6 +29,12 @@ namespace sfm {
struct MatchOptions {
float max_ratio = 0.8f; // Lowe ratio (best_dist < ratio * second_dist)
// Absolute cosine similarity a match must reach, for L2-normalized float
// descriptors only; 0 disables it. COLMAP's ALIKED defaults are
// min_cossim = 0.85 with max_ratio = 1.0 -- i.e. the ratio test off and
// this the only filter -- which is a different shape of test from SIFT's
// and needs both knobs to exist to be expressible.
float min_similarity = 0.0f;
bool cross_check = true; // keep only mutual nearest neighbours
uint32_t max_num_matches = 32768; // cap per pair (0 = unlimited)
int device = -1;
@@ -288,6 +294,35 @@ private:
}
}
// Float descriptors reach the GPU as uint8.
//
// The matching kernel is built around the hardware packed uint8x4 dot
// product and a 128-byte descriptor; a float path would quadruple its
// groupshared tile and retune it. It does not have to: for a UNIFORM
// affine quantization the squared distances all scale by one constant, so
// the ordering, the Lowe ratio (a ratio of two of them) and the
// max_num_matches sort are preserved exactly. Only an absolute threshold
// has to be converted, which is what similarityFromD2 does.
//
// kQuantHalfRange is the half-range mapped onto the byte. Measured on real
// ALIKED descriptors: components are ~N(0, 0.088) with a maximum of 0.392
// over an image, so 0.4 clips essentially nothing and spends the byte on
// the range that is actually occupied. The residual is ~1% of a unit
// descriptor's norm, well under the ratio test's margin.
static constexpr float kQuantHalfRange = 0.4f;
static constexpr float kQuantSteps = 127.0f;
static uint8_t quantize(float v) {
const float q = v * (kQuantSteps / kQuantHalfRange) + 128.0f;
return (uint8_t)std::lround(std::min(255.0f, std::max(0.0f, q)));
}
// Cosine similarity of two unit descriptors from their quantized squared
// distance: ||a-b||^2 = 2 - 2 cos, and d2 is that scaled by (steps/range)^2.
static float similarityFromD2(uint32_t d2) {
const float k = kQuantHalfRange / kQuantSteps;
return 1.0f - 0.5f * (float)d2 * k * k;
}
static constexpr int DW = 32; // uint words per descriptor
// Workgroup width of the row-only kernel; MUST equal TQR in
// sfm/shaders/match/bruteforce.slang. Each workgroup streams the whole
@@ -297,10 +332,15 @@ private:
// queries that are not even in range.
static constexpr uint32_t kRowThreads = 256;
// 128 bytes per descriptor either way: uint8 as SIFT writes them, or f32
// quantized on upload (see kQuantHalfRange).
static void checkDescriptors(const FeatureSet& f) {
if (f.count() == 0) return;
if (f.dtype != DType::U8 || f.dim != 128)
throw std::runtime_error("brute-force matcher expects 128-D uint8 descriptors");
if (f.dim != 128)
throw std::runtime_error("brute-force matcher expects 128-D descriptors, got " +
std::to_string(f.dim));
if (f.dtype != DType::U8 && f.dtype != DType::F32)
throw std::runtime_error("brute-force matcher expects uint8 or f32 descriptors");
}
// Lowe ratio + cross-check, exactly as before -- the GPU only changed how
@@ -314,6 +354,8 @@ private:
uint32_t bestD2 = rA[4 * i + 1], secondD2 = rA[4 * i + 2];
if (j >= nb) continue;
if (secondD2 != 0xffffffffu && (float)bestD2 >= r2 * (float)secondD2) continue;
if (opt_.min_similarity > 0 && similarityFromD2(bestD2) < opt_.min_similarity)
continue;
if (opt_.cross_check && rB[4 * j + 0] != i) continue;
out.push_back({i, j, std::sqrt((float)bestD2)});
}
@@ -338,6 +380,8 @@ private:
uint32_t bestD2 = rA[4 * i + 1], secondD2 = rA[4 * i + 2];
if (j >= nb) continue;
if (secondD2 != 0xffffffffu && (float)bestD2 >= r2 * (float)secondD2) continue;
if (opt_.min_similarity > 0 && similarityFromD2(bestD2) < opt_.min_similarity)
continue;
if (opt_.cross_check && rB[4 * j + 0] != i) continue;
n++;
}
@@ -442,10 +486,15 @@ private:
const FeatureSet& f = *feats[img];
resident_[img] = used_;
used_ += f.count();
if (f.count()) {
memcpy(blob + off, f.descriptors.data(), (size_t)f.count() * 128);
off += (size_t)f.count() * 128;
if (!f.count()) continue;
const size_t bytes = (size_t)f.count() * 128;
if (f.dtype == DType::U8) {
memcpy(blob + off, f.descriptors.data(), bytes);
} else {
const float* src = reinterpret_cast<const float*>(f.descriptors.data());
for (size_t i = 0; i < bytes; i++) blob[off + i] = quantize(src[i]);
}
off += bytes;
}
ctx_.upload(bDesc_, blob, (VkDeviceSize)want * 128, (VkDeviceSize)first * 128);
normDirty_ = {first, used_};
src/sfm/feature/PairSelection.h
+11 -6
View File
@@ -71,17 +71,22 @@ struct PairSelectionOptions {
int device = -1;
};
// Gather f's K largest-scale features (K = 0 or >= count keeps everything,
// but still gathered in scale order). The canonical feature order is by
// position, deliberately not scale (D16), so this is an explicit host-side
// gather; scale ties break by index.
// Gather f's K best-ranked features (K = 0 or >= count keeps everything, but
// still gathered in rank order). The canonical feature order is by position,
// deliberately not by rank (D16), so this is an explicit host-side gather;
// ties break by index.
//
// "Rank" is scale for SIFT and the detection score for a detector that has no
// scale -- FeatureSet::rank picks. Reading `scale` unconditionally, as this
// did, silently selected an arbitrary 512 keypoints out of an ALIKED set,
// because every one of them has scale 0. That is the failure mode this whole
// stage is least able to show: the shortlist still looks populated.
inline FeatureSet topScaleSubset(const FeatureSet& f, uint32_t K) {
if (K == 0 || K > f.count()) K = f.count();
std::vector<uint32_t> idx(f.count());
for (uint32_t j = 0; j < f.count(); j++) idx[j] = j;
std::partial_sort(idx.begin(), idx.begin() + K, idx.end(), [&](uint32_t a, uint32_t b) {
if (f.keypoints[a].scale != f.keypoints[b].scale)
return f.keypoints[a].scale > f.keypoints[b].scale;
if (f.rank(a) != f.rank(b)) return f.rank(a) > f.rank(b);
return a < b;
});
FeatureSet m;
tools/aliked/compare_colmap.py
+130
View File
@@ -0,0 +1,130 @@
#!/usr/bin/env python3
"""Compare our ALIKED extractor against COLMAP's ONNX one, on the same image.
This is the gate that says the port reproduces the reference, and it is the
reason the checkpoint is fetched from COLMAP's releases rather than converted:
both sides run the same weights, so a difference here is ours.
# 1. COLMAP's side (needs a COLMAP built with ONNX support)
colmap feature_extractor --database_path /tmp/db.db --image_path IMAGES \\
--FeatureExtraction.type ALIKED_N16ROT
# 2. ours
./build/aliked_test --image IMAGES/x.jpg --out /tmp/ours.bin
# 3. compare
python3 tools/aliked/compare_colmap.py /tmp/db.db /tmp/ours.bin
What it reports, and what each number means:
matched the fraction of our keypoints that have a COLMAP keypoint within
`--tol` pixels. Detection is a threshold on a continuous score
map, so a handful of borderline points differ between any two
implementations; what would signal a real bug is a systematic
miss, or an offset in the mean residual.
offset the mean signed (dx, dy) over matched pairs. This is the number
that catches a half-pixel convention error, and it should be
zero to several decimals, not merely small.
cosine descriptor agreement on matched pairs. Both are L2-normalized,
so this is a dot product; anything below ~0.99 means the
descriptor head diverges even where the detector agrees.
"""
import argparse
import sqlite3
import struct
import sys
import numpy as np
def read_ours(path):
with open(path, "rb") as f:
blob = f.read()
if blob[:8] != b"ALIKEDFT":
raise SystemExit(f"{path}: not an aliked_test dump")
version, w, h, count, dim = struct.unpack_from("<IiiII", blob, 8)
if version != 1:
raise SystemExit(f"{path}: version {version}")
off = 28
kp = np.frombuffer(blob, np.float32, count * 3, off).reshape(count, 3)
off += count * 3 * 4
desc = np.frombuffer(blob, np.float32, count * dim, off).reshape(count, dim)
return dict(width=w, height=h, xy=kp[:, :2], score=kp[:, 2], desc=desc)
def read_colmap(db_path, image_name=None):
con = sqlite3.connect(db_path)
rows = con.execute("SELECT image_id, name FROM images ORDER BY image_id").fetchall()
if not rows:
raise SystemExit(f"{db_path}: no images")
if image_name is not None:
rows = [r for r in rows if r[1] == image_name or r[1].endswith("/" + image_name)]
if not rows:
raise SystemExit(f"{db_path}: no image named {image_name}")
image_id, name = rows[0]
r, c, data = con.execute(
"SELECT rows, cols, data FROM keypoints WHERE image_id=?", (image_id,)
).fetchone()
kp = np.frombuffer(data, np.float32).reshape(r, c)
r2, c2, data2 = con.execute(
"SELECT rows, cols, data FROM descriptors WHERE image_id=?", (image_id,)
).fetchone()
# ALIKED descriptors are float32 stored in COLMAP's uint8 descriptor blob
# (feature/aliked.cc writes descriptor_dim * sizeof(float) bytes per row),
# so the column count is 4x the descriptor width.
desc = np.frombuffer(data2, np.uint8).reshape(r2, c2).view(np.float32)
con.close()
return dict(name=name, xy=kp[:, :2].copy(), desc=desc)
def main():
ap = argparse.ArgumentParser()
ap.add_argument("database")
ap.add_argument("ours")
ap.add_argument("--image", default=None, help="image name in the database")
ap.add_argument("--tol", type=float, default=1.0, help="match radius, pixels")
args = ap.parse_args()
ours = read_ours(args.ours)
ref = read_colmap(args.database, args.image)
print(f"ours : {len(ours['xy']):5d} keypoints, {ours['desc'].shape[1]}-D")
print(f"colmap : {len(ref['xy']):5d} keypoints, {ref['desc'].shape[1]}-D ({ref['name']})")
if ours["desc"].shape[1] != ref["desc"].shape[1]:
raise SystemExit("descriptor widths differ")
# Nearest reference keypoint for each of ours. A few thousand points each
# way, so the quadratic distance matrix is the simple thing that works.
d = np.linalg.norm(ours["xy"][:, None, :] - ref["xy"][None, :, :], axis=2)
nearest = d.argmin(axis=1)
dist = d[np.arange(len(nearest)), nearest]
ok = dist <= args.tol
n_ok = int(ok.sum())
print(f"\nmatched: {n_ok}/{len(ok)} ({100.0 * n_ok / max(len(ok), 1):.1f}%) "
f"within {args.tol} px")
if n_ok == 0:
raise SystemExit("no keypoints matched -- check the coordinate convention")
delta = ref["xy"][nearest[ok]] - ours["xy"][ok]
print(f"offset : mean ({delta[:, 0].mean():+.4f}, {delta[:, 1].mean():+.4f}) px, "
f"rms {np.sqrt((delta ** 2).sum(axis=1).mean()):.4f} px")
print(f" max |dx| {np.abs(delta[:, 0]).max():.4f}, "
f"max |dy| {np.abs(delta[:, 1]).max():.4f}")
cos = (ours["desc"][ok] * ref["desc"][nearest[ok]]).sum(axis=1)
print(f"cosine : mean {cos.mean():.6f}, min {cos.min():.6f}, "
f"5th pct {np.percentile(cos, 5):.6f}")
below = int((cos < 0.99).sum())
print(f" {below} of {n_ok} below 0.99")
# Not a pass/fail: what "close enough" means depends on what is being
# changed, and a human reading the three numbers above is the gate.
return 0
if __name__ == "__main__":
sys.exit(main())