Python:如何通过 TensorRT-Model-Connect 部署 LocateAnything-3B?

作者:袖梨 2026-09-17

TensorRT-Model-Connect 已提供 LocateAnything-3B 的官方 family recipe,可以把 Hugging Face 检查点直接构建为包含视觉、prefill 和 decode TensorRT engine 的 .bundle,再通过原生 trtmc 运行图像定位。当前配方适合评估和参考实现:只支持 FP16、单卡、batch size 1、非量化构建,图像尺寸固定为 448×448,不能开启动态 KV cache。

先确认项目状态与平台边界

TensorRT-Model-Connect 当前标注为 Public Preview,API、支持范围和构建方式仍可能变化。发布 wheel 的路径目前只覆盖 Linux aarch64、Python 3.10 或 3.12、glibc 2.39 及以上,并匹配官方 TensorRT 11.1.0.106。Linux x86_64 目前应使用源码构建路径,且需要 Docker、NVIDIA Container Toolkit 和可用的 NVIDIA GPU。

不要把不同架构、TensorRT 版本批次或安装来源的组件混在一起。bundle 构建时生成的 TensorRT plan 与 GPU/TensorRT 环境有关,运行时使用的 libtrtmc_core.solibtrtmc_runtime.so、后端 DSO 和 LocateAnything family DSO 也必须来自匹配的构建。

理解 LocateAnything recipe 的硬约束

官方模型配方声明的 checkpoint 是 nvidia/LocateAnything-3B,任务接口为 vision_language_generation,默认和测试精度为 FP16,tensor parallel size 为 1。family 构建器还会拒绝以下请求:

  • 动态 KV cache。
  • 自定义 image_heightimage_width
  • batch size 不等于 1。
  • context parallel 或 tensor parallel 大于 1。
  • 量化、指定 FP32 层、视频帧输入或其他任务类型。

视觉链路固定使用 448×448 输入,经 MoonViT 形成视觉 token。官方测试 manifest 使用 max_sequence_length=384,包含单目标框定位和点定位两种用例。先复现这个范围,再尝试修改 recipe,能把环境问题与模型扩展问题分开。

从源码准备构建环境

以下流程适合 Linux x86_64 或 aarch64。先取得 TensorRT-Model-Connect 源码,并确保 Docker 能通过 NVIDIA Container Toolkit 访问 GPU。项目提供开发 Dockerfile,可根据主机架构选择对应文件。

git clone <TensorRT-Model-Connect仓库克隆地址>
cd TensorRT-Model-Connect

GPU=0
SM="$(
  nvidia-smi -i "$GPU" 
    --query-gpu=compute_cap 
    --format=csv,noheader,nounits | tr -d '.[:space:]'
)"

case "$(uname -m)" in
  x86_64) DOCKERFILE=Dockerfile.dev.x86 ;;
  aarch64) DOCKERFILE=Dockerfile.dev.aarch64 ;;
  *) echo "不支持的主机架构"; exit 1 ;;
esac

docker build 
  -f "$DOCKERFILE" 
  -t trtmc-locateanything 
  requirements

SM 是所选 GPU 的 compute capability 去掉小数点后的值,例如 8.9 会转成 89。该值稍后传给 CMake,避免构建与目标 GPU 不匹配的 CUDA 架构。

启动带 GPU 的开发容器

把当前源码目录挂载到容器的 /src,使用 --ipc=host,并只暴露选定 GPU:

SOURCE_DIR="$(git rev-parse --show-toplevel)"

docker run --rm -it 
  --gpus "device=${GPU}" 
  --ipc=host 
  --mount "type=bind,source=${SOURCE_DIR},target=/src" 
  --workdir /src 
  --env TRTMC_SM="$SM" 
  trtmc-locateanything 
  bash

后面的 Python 安装、CMake 编译、bundle 构建和运行命令都在这个容器 shell 中执行。若 nvidia-smi 在容器内不可见,应先修复 Docker GPU 访问,而不是继续排查 TensorRT 模型代码。

安装 Python 构建端与 family 依赖

Python 模块负责读取 Hugging Face 检查点并构建 bundle,原生 CLI 负责检查和运行 bundle。先以只安装 Python 部分的方式安装当前源码,再安装 LocateAnything family 的额外依赖:

python -m pip install --no-deps -e . -C py-only=true
python -m pip install -r families/locateanything/requirements.txt

当前 family requirements 包含 Pillow、PEFT,并固定 transformers==4.57.6。不要用全局环境中另一套 Transformers 覆盖该约束;构建器需要按当前仓库预期读取模型配置、tokenizer 和权重。

编译原生运行时和 LocateAnything DSO

与通用快速开始中的 Qwen 示例不同,这里必须编译 trtmc_model_locateanything,否则 bundle 虽能构建,运行时也会因为缺少 family DSO 而加载失败。

TRTMC_BUILD_DIR="build-sm${TRTMC_SM}"

cmake -S . -B "$TRTMC_BUILD_DIR" -G Ninja 
  -DCMAKE_BUILD_TYPE=Release 
  -DCMAKE_CUDA_ARCHITECTURES="${TRTMC_SM}-real" 
  -DTRTMC_BUILD_BACKEND_RTX=OFF 
  -DTRTMC_BUILD_TESTS=OFF 
  -DTRTMC_BUILD_EXAMPLES=OFF

cmake --build "$TRTMC_BUILD_DIR" 
  --parallel "$(nproc)" 
  --target 
    trtmc 
    trtmc_backend_trt 
    trtmc_model_locateanything

export PATH="$PWD/$TRTMC_BUILD_DIR:$PATH"
TRTMC_RUNTIME_ROOT="$PWD/$TRTMC_BUILD_DIR"

运行目录至少应包含核心运行库、TensorRT 后端库和 libtrtmc_model_locateanything.sotrtmc 不会自动搜索当前目录、环境变量中的备用目录或其他安装位置,因此每次运行都要显式传入 --runtime-root

构建 LocateAnything-3B bundle

构建命令可以直接使用 Hugging Face 模型 ID,也可以传入已经准备好的本地 snapshot 目录。为了贴近官方 manifest,显式指定 FP16、384 序列长度、单 batch 和单卡:

python -m tensorrt_model_connect build 
  nvidia/LocateAnything-3B 
  --task vision_language_generation 
  --precision fp16 
  --max-sequence-length 384 
  --max-batch-size 1 
  --tensor-parallel-size 1 
  --output locateanything-3b.bundle 
  --verbose

该过程会下载或读取模型 snapshot,加载 LocateAnything family builder,分别构建 decode engine、prefill engine 和 vision engine,并把运行时配置与 tokenizer 资源写进 bundle。构建可能消耗较长时间和大量磁盘、主机内存及显存;不要因为终端一段时间没有新日志就启动第二个相同构建。

配方当前不支持量化,因此不要添加 --quantization。也不要传入图片宽高参数,固定 448×448 的预处理契约由 family 自己写入运行配置。

检查 bundle 元数据

构建成功后先用 inspect 检查文件头,而不是直接运行:

trtmc inspect locateanything-3b.bundle
trtmc version

检查结果应能识别 family 为 locateanything、task 为 vision_language_generation、backend 为标准 TensorRT,并列出 bundle 分区范围。inspect 只证明容器格式可读,不代表模型输出已经通过质量验证。

运行单图框定位

原生任务使用 run,图像通过 --image 提供,提示词通过 --prompt 提供。下面沿用官方测试用例的单目标模板和 32 个新 token 上限:

trtmc run locateanything-3b.bundle 
  --runtime-root "$TRTMC_RUNTIME_ROOT" 
  --image test_img.jpeg 
  --prompt "Locate a single instance that matches the following description: white vehicle." 
  --max-new-tokens 32 
  --temperature 1 
  --top-k 1

temperature=1 配合 top-k=1 最终仍只选择概率最高的候选,与官方 E2E 调用保持一致。输出应包含 JSON,其中 text 字段携带结构化定位 token。不要在没有实际执行的情况下预写具体坐标,因为坐标取决于输入图片和模型输出。

运行点定位

同一个 bundle 也可通过修改提示词执行点定位。官方测试使用如下模板:

trtmc run locateanything-3b.bundle 
  --runtime-root "$TRTMC_RUNTIME_ROOT" 
  --image test_img.jpeg 
  --prompt "Point to: white vehicle." 
  --max-new-tokens 32 
  --temperature 1 
  --top-k 1

框定位与点定位都属于同一个视觉语言生成任务,并不需要分别构建 engine。应用程序应解析模型输出中的坐标 token,再依据模型的归一化坐标约定还原到原图尺寸。

常见构建失败如何定位

family 未识别模型

构建器会读取 checkpoint 根目录的 config.json,并要求恰好一个 family 声明支持。LocateAnything family 只接受 model_type=locateanything。若传入的是不完整下载目录、错误模型或被修改的配置,会在 TensorRT engine 构建前失败。

请求了 recipe 不支持的选项

出现 dynamic KV、图片尺寸、batch、并行或量化相关的 NotImplementedError 时,应删除这些参数,回到单卡 FP16 基线。它不是显存自动降级提示,当前 family 不会在不支持的参数上静默回退。

运行时提示缺少 DSO

确认 --runtime-root 指向本次 CMake 构建目录,并检查核心库、标准 TensorRT 后端和 LocateAnything family 库同时存在。不要把 wheel 中的 CLI 与另一个源码构建目录下的 DSO 混用。

TensorRT plan 与设备不兼容

核对 CMake 的 CUDA architecture 是否来自实际运行 GPU,并在目标环境重新构建 bundle。把一个环境生成的 plan 复制到不同 TensorRT 批次或不兼容 GPU 上,不能视为可移植部署方式。

部署前应验证什么

至少保留构建仓库提交、模型 revision、GPU 型号、compute capability、TensorRT 版本、精度、序列长度和 bundle 哈希。使用固定图片分别运行框定位与点定位,并保存原生 JSON 输出;随后与官方 Transformers 参考实现对比文本结构和坐标误差。

TensorRT-Model-Connect 在这里承担的是“从受支持 checkpoint 到原生任务接口”的参考路径,不是自动扩展到多卡、量化和任意图片尺寸的生产平台。先在官方约束内构建 FP16 单卡 bundle,验证视觉预处理、tokenizer、prefill、decode 和坐标输出都一致,再根据项目源码评估需要扩展的边界,风险最低。

相关文章

精彩推荐