DLA¶
DLA
NVIDIA 深度學習加速器是針對深度學習運算而設計的固定功能加速器引擎。DLA 旨在對卷積神經網路進行完整的硬體加速。DLA 支援各種層,例如卷積、反卷積、全連接、啟動、池化、批次正規化等。torch_tensorrt
支援在 NVIDIA 嵌入式平台上可用的 DLA 硬體上編譯 TorchScript 模組和部署管線。
注意:DLA 僅支援 fp16 和 int8 精度。
搭配 torchtrtc 使用 DLA
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)