dtourolle/scene-actor-extraction
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Add AMD support via ort alternative to trt

Browse Source
This commit is contained in:
dtourolle
2026-06-28 11:50:05 +02:00
parent a3ba53ddf7
commit 0ee131a692
27 changed files with 1357 additions and 977 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/nodes/identity_matcher_node.hpp
+10 -88
View File
@@ -1,16 +1,15 @@
#pragma once
#include "types.hpp"
#include "config.hpp"
#include "inference/similarity.hpp"
#include "gallery/gallery_store.hpp"
#include "gallery/gallery_calibration.hpp"
#include <cublas_v2.h>
#include <cuda_runtime_api.h>
#include <cstdint>
#include <cstring>
#include <iostream>
#include <limits>
#include <memory>
#include <stdexcept>
#include <vector>
@@ -35,25 +34,9 @@
// Gallery scan: the full reference set (tens of thousands of 512-dim
// embeddings) is uploaded to the GPU once at construction time and stays
// resident there. Per frame, only the small query matrix (n_faces x 512) is
// uploaded and a single cublasSgemm computes the full similarity matrix
// (n_faces x N_gallery) in well under a millisecond — far faster than any
// CPU GEMM or scalar scan, making gallery-side pruning unnecessary.
namespace identity_matcher_detail {
struct CudaError : std::runtime_error {
using std::runtime_error::runtime_error;
};
inline void check_cuda(cudaError_t e, const char* what) {
if (e != cudaSuccess)
throw CudaError(std::string(what) + ": " + cudaGetErrorString(e));
}
inline void check_cublas(cublasStatus_t s, const char* what) {
if (s != CUBLAS_STATUS_SUCCESS)
throw CudaError(std::string(what) + ": cublas error " + std::to_string(s));
}
} // namespace identity_matcher_detail
// uploaded and a single SGEMM computes the full similarity matrix in well under
// a millisecond. The GPU math backend (cuBLAS or rocBLAS) lives behind
// ISimilarityEngine (backends/gemm_backend.cpp) and is selected at compile time.
struct IdentityMatcherFunc {
static constexpr std::string_view label() { return "identity_matcher"; }
@@ -69,8 +52,6 @@ struct IdentityMatcherFunc {
, ratio_(cfg.match_ratio)
, ratio_ceil_(cfg.match_ratio_ceil)
{
using namespace identity_matcher_detail;
std::cerr << "[identity_matcher] flattening gallery embeddings...\n";
for (int ai = 0; ai < static_cast<int>(gallery_.actors.size()); ++ai) {
for (const auto& emb : gallery_.actors[ai].embeddings) {
@@ -100,47 +81,15 @@ struct IdentityMatcherFunc {
<< gallery_.actors.size() << " actors, "
<< flat_emb_.size() << " reference embeddings\n";
// Flatten gallery into a contiguous (N x 512) row-major host buffer,
// then upload once. Row-major NxD == column-major DxN, which is the
// layout cublasSgemm wants for the transposed operand below.
std::vector<float> host_gallery(static_cast<size_t>(n_gallery_) * 512);
for (int i = 0; i < n_gallery_; ++i)
std::memcpy(host_gallery.data() + static_cast<size_t>(i) * 512,
flat_emb_[i].data(), 512 * sizeof(float));
check_cuda(cudaMalloc(reinterpret_cast<void**>(&d_gallery_), host_gallery.size() * sizeof(float)),
"cudaMalloc gallery");
check_cuda(cudaMemcpy(d_gallery_, host_gallery.data(),
host_gallery.size() * sizeof(float),
cudaMemcpyHostToDevice),
"cudaMemcpy gallery H2D");
check_cuda(cudaMalloc(reinterpret_cast<void**>(&d_query_), static_cast<size_t>(kMaxFaces) * 512 * sizeof(float)),
"cudaMalloc query");
check_cuda(cudaMalloc(reinterpret_cast<void**>(&d_sims_), static_cast<size_t>(kMaxFaces) * n_gallery_ * sizeof(float)),
"cudaMalloc sims");
check_cuda(cudaStreamCreate(&stream_), "cudaStreamCreate");
check_cublas(cublasCreate(&handle_), "cublasCreate");
check_cublas(cublasSetStream(handle_, stream_), "cublasSetStream");
host_sims_.resize(static_cast<size_t>(kMaxFaces) * n_gallery_);
std::cerr << "[identity_matcher] gallery resident on GPU ("
<< (host_gallery.size() * sizeof(float)) / (1024 * 1024) << " MiB)\n";
}
~IdentityMatcherFunc() {
if (d_gallery_) cudaFree(d_gallery_);
if (d_query_) cudaFree(d_query_);
if (d_sims_) cudaFree(d_sims_);
if (handle_) cublasDestroy(handle_);
if (stream_) cudaStreamDestroy(stream_);
sim_engine_ = make_similarity_engine(host_gallery.data(), n_gallery_, kMaxFaces);
}
MatchedSceneFrame operator()(TrackedSceneFrame tf) {
using namespace identity_matcher_detail;
if (tf.source.eof) return {std::move(tf.source), {}};
const int n_faces = static_cast<int>(tf.embeddings.size());
@@ -151,37 +100,18 @@ struct IdentityMatcherFunc {
if (n_faces > kMaxFaces)
throw std::runtime_error("identity_matcher: n_faces exceeds kMaxFaces");
// Build the (n_faces x 512) query matrix (row-major == col-major 512 x n_faces).
std::vector<float> host_query(static_cast<size_t>(n_faces) * 512);
for (int fi = 0; fi < n_faces; ++fi) {
std::memcpy(host_query.data() + static_cast<size_t>(fi) * 512,
tf.embeddings[fi].data(), 512 * sizeof(float));
}
check_cuda(cudaMemcpyAsync(d_query_, host_query.data(),
host_query.size() * sizeof(float),
cudaMemcpyHostToDevice, stream_),
"cudaMemcpy query H2D");
// S (N_gallery x n_faces) col-major = G(512 x N_gallery)^T * Q(512 x n_faces)
// i.e. S[g + f*N_gallery] = cosine_similarity(gallery[g], query[f]).
const float alpha = 1.f, beta = 0.f;
check_cublas(cublasSgemm(handle_, CUBLAS_OP_T, CUBLAS_OP_N,
n_gallery_, n_faces, 512,
&alpha, d_gallery_, 512, d_query_, 512,
&beta, d_sims_, n_gallery_),
"cublasSgemm");
check_cuda(cudaMemcpyAsync(host_sims_.data(), d_sims_,
static_cast<size_t>(n_gallery_) * n_faces * sizeof(float),
cudaMemcpyDeviceToHost, stream_),
"cudaMemcpy sims D2H");
check_cuda(cudaStreamSynchronize(stream_), "cudaStreamSynchronize");
// S (N_gallery × n_faces) col-major: face fi's gallery sims at sims + fi*n_gallery.
const float* host_sims = sim_engine_->compute(host_query.data(), n_faces);
for (int fi = 0; fi < n_faces; ++fi) {
const float* sims = host_sims_.data() + static_cast<size_t>(fi) * n_gallery_;
const float* sims = host_sims + static_cast<size_t>(fi) * n_gallery_;
// Per-actor best cosine similarity (max over that actor's reference embeddings)
std::vector<float> best_sim(gallery_.actors.size(),
-std::numeric_limits<float>::max());
for (int ei = 0; ei < n_gallery_; ++ei) {
@@ -190,7 +120,6 @@ struct IdentityMatcherFunc {
if (sim > best_sim[ai]) best_sim[ai] = sim;
}
// Find best and second-best actor by similarity
int best_actor = -1;
int second_actor = -1;
float best_s = -std::numeric_limits<float>::max();
@@ -240,7 +169,6 @@ struct IdentityMatcherFunc {
? cal_.probability(best_s, log_prior_odds_)
: best_s;
}
// actor_idx == -1, name == "" → unknown face
actors.push_back(std::move(ia));
}
@@ -260,11 +188,5 @@ private:
std::vector<int> flat_actor_;
int n_gallery_{0};
float* d_gallery_{nullptr}; // (n_gallery_ x 512), resident for the lifetime of this node
float* d_query_{nullptr}; // (kMaxFaces x 512)
float* d_sims_{nullptr}; // (n_gallery_ x kMaxFaces), column-major
std::vector<float> host_sims_;
cudaStream_t stream_{nullptr};
cublasHandle_t handle_{nullptr};
std::unique_ptr<ISimilarityEngine> sim_engine_;
};