DLA

DLA NVIDIA 深度学习加速器是一种固定功能加速器引擎，专用于深度学习操作。DLA 设计用于对卷积神经网络进行完全硬件加速。DLA 支持各种层，如卷积、反卷积、全连接、激活、池化、批归一化等。torch_tensorrt 支持在 NVIDIA 嵌入式平台上的 DLA 硬件上编译 TorchScript 模块和部署流水线。

注意：DLA 仅支持 fp16 和 int8 精度。

将 DLA 与 torchtrtc 一起使用

torchtrtc [input_file_path] [output_file_path] [input_shapes...] -p f16 -d dla {OPTIONS}

在 C++ 应用中使用 DLA

std::vector<std::vector<int64_t>> input_shape = {{32, 3, 32, 32}};
auto compile_spec = torch_tensorrt::CompileSpec({input_shape});

# Set a precision. DLA supports fp16 or int8 only
compile_spec.enabled_precisions = {torch::kF16};
compile_spec.device.device_type = torch_tensorrt::CompileSpec::DeviceType::kDLA;

# Make sure the gpu id is set to Xavier id for DLA
compile_spec.device.gpu_id = 0;

# Set the DLA core id
compile_spec.device.dla_core = 1;

# If a layer fails to run on DLA it will fallback to GPU
compile_spec.device.allow_gpu_fallback = true;

在 Python 应用中使用 DLA

compile_spec = {
    "inputs": [torch_tensorrt.Input(self.input.shape)],
    "device": torch_tensorrt.Device("dla:0", allow_gpu_fallback=True),
    "enabled_precisions": {torch.half},
}

trt_mod = torch_tensorrt.compile(self.scripted_model, compile_spec)