pytorch / TensorRT

PyTorch/TorchScript/FX compiler for NVIDIA GPUs using TensorRT
https://pytorch.org/TensorRT
BSD 3-Clause "New" or "Revised" License
2.54k stars 349 forks source link

🐛 [Bug] Outdated "Compiling with Torch-TensorRT in C++" Documentation #2947

Closed Borotalcohol closed 2 months ago

Borotalcohol commented 3 months ago

Bug Description

In the "Using Torch-TensorRT in C++" page, under the "Compiling with Torch-TensorRT in C++" section, the syntax of an older version of the library is being used:

Example

#include "torch/script.h"
#include "torch_tensorrt/torch_tensorrt.h"
...

mod.to(at::kCUDA);
mod.eval();

auto in = torch::randn({1, 1, 32, 32}, {torch::kCUDA});
auto trt_mod = torch_tensorrt::CompileGraph(mod, std::vector<torch_tensorrt::CompileSpec::InputRange>{{in.sizes()}});
auto out = trt_mod.forward({in});

torch_tensorrt doesn't have a CompileGraph method and torch_tensorrt::CompileSpec, moved now to torch_tensorrt::ts::CompileSpec, doesn't have an InputRange class.

The same result should be achievable right now with something like:

#include "torch/script.h"
#include "torch_tensorrt/torch_tensorrt.h"
...

mod.to(at::kCUDA);
mod.eval();

auto in = torch::randn({1, 1, 32, 32}, {torch::kCUDA});
torch_tensorrt::ts::CompileSpec spec({1, 1, 32, 32});
std::vector<torch_tensorrt::Input> inputs;
inputs.push_back(torch_tensorrt::Input(in));
spec.graph_inputs.inputs = inputs;
auto trt_mod = torch_tensorrt::ts::compile(mod, spec);
auto out = trt_mod.forward({in});

Note, however, that syntax should be updated also in subsequent code snippets.