Is TensorRT hiding your AI model weights?

This article aims to show that a compiled TensorRT engine doesn’t encrypt its weights, and that the weights are hard-coded and easy to extract with small example.

What is TensorRT?

NVIDIA TensorRT is an SDK (Software Development Kit) designed for high-performance inferences. While frameworks like TensorFlow or PyTorch are used to train models, TensorRT is used to optimize (via compilation) and deploy them on NVIDIA GPUs to make them run as fast as possible, for exemple on NVIDIA RTX or NVIDIA JETSON.

Compilation is the process of transforming source code into a format that a computer’s processor or a specific hardware can execute. In the TensorRT context, compilation gives a standardize format of the AI model to the GPU and what to do with it through a new file to be executed: a .engine file.

Since TensorRT runs inferences, he needs the weights at some point, NVIDIA doesn’t handle encryption during compilation, so the weights must be hard-coded in the engine file.

Figure 1: TensorRT engine file structure in hexadecimal.

In our previous article, we found and used this structure (Figure 1) to extract the tensor sizes and to create signatures of models (to be more precise, of their architecture). But we can also just extract directly the weights of the model, this is the purpose of this article.

On what device ?

Since the NVIDIA Jetson ORIN is one of the most popular device that uses TensorRT, we will perform the demonstration on it.

The NVIDIA Jetson is a series of efficient computing modules made for edge AI and embedded systems. Built around NVIDIAs GPU architecture, Jetson platforms give access to a significant computing power on a small form factor, enabling real time inference, computer vision, and deep learning tasks directly on the device.

Its popularity comes from a few key strengths: it removes the gap between powerful data center GPUs and the constraints of embedded hardware, making it ideal for robotics, smart cameras etc. Developers also benefit from NVIDIAs mature ecosystem (CUDA, TensorRT), which makes deploying AI models easier.

Weights extraction

To demonstrate how easy it is, here is an example, we created a small neural network to extract a weights batch of the first layer. We used the same code as the one of ModelDNA, expect instead of counting layer size weights by weights (in hexadecimal, they are float32), we just print them.

Figure 2: Demonstration of how weights are hard-coded into the engine file.

We find the weights as expected, they are hard-coded and even in the right order. You can imagine that this is replicable to any layers.

Is your model vulnerable ?

If you are using TensorRT compilation as a protection, yes, if an attacker uses our ModelDNA to find the architecture of your model, he can easily try to dump the weights into a blank model to get an exact copy of yours.

The AI protection of Skyld helps you to protect the parameters of an AI model (weights) against extraction before and during the model execution, using mathematical transformations to keep the model weights confidential and thus make obsolete this type of attack.