■PyTorch → ONNX → TensorRT 手順書

雪の白川郷

① 必要なもの（最初に全部）
CUDA 12.1 の場合
pip install torch torchvision torchaudio –index-url https://download.pytorch.org/whl/cu121
torch
onnx
onnxscript
TensorRT（trtexec）
Docker（NVIDIA GPU 対応）

② PyTorch → ONNX（5 行）
import torch
import torch.nn as nn
model = nn.Linear(4, 2)
dummy = torch.randn(1, 4)
torch.onnx.export(model, dummy, “model.onnx”)

● 保存場所
Python を起動した場所（例：/home/chibi）

● 出力
model.onnx
model.onnx.data

③ Docker に入る（RHEL10 / Ubuntu 共通）
sudo docker run –rm –gpus all -it \
-v $HOME:/workspace \
nvcr.io/nvidia/tensorrt:24.02-py3

④ ONNX → TensorRT（1 行）
trtexec –onnx=/workspace/model.onnx –saveEngine=/workspace/model.plan
ls -al /workspace

● 出力
model.plan

⑤ 推論（必要なら）
trtexec –loadEngine=/workspace/model.plan –shapes=input:1×4

データをコピペで使う場合はこちらでお願いします。
データに使うコピペ版