torch2trt
v0.5.0
คุณใช้รุ่นใดหรือหวังว่าจะใช้กับ Tensorrt? อย่าลังเลที่จะเข้าร่วมการสนทนาที่นี่
Torch2trt เป็น Pytorch to Tensorrt Converter ซึ่งใช้ Tensorrt Python API ตัวแปลงคือ
ใช้งานง่าย - แปลงโมดูลด้วยฟังก์ชั่นเดียวเรียก torch2trt
ง่ายต่อการขยาย - เขียนตัวแปลงเลเยอร์ของคุณเองใน Python และลงทะเบียนด้วย @tensorrt_converter
หากคุณพบปัญหาโปรดแจ้งให้เราทราบ!
โปรดทราบว่าตัวแปลงนี้มีความครอบคลุม จำกัด ของ Tensorrt / Pytorch เราสร้างมันขึ้นมาเพื่อเพิ่มประสิทธิภาพโมเดลที่ใช้ในโครงการ JetBot ได้อย่างง่ายดาย หากคุณพบว่าตัวแปลงมีประโยชน์กับรุ่นอื่น ๆ โปรดแจ้งให้เราทราบ
ด้านล่างนี้เป็นตัวอย่างการใช้งานบางส่วนสำหรับการตรวจสอบเพิ่มเติมสมุดบันทึก
import torch
from torch2trt import torch2trt
from torchvision . models . alexnet import alexnet
# create some regular pytorch model...
model = alexnet ( pretrained = True ). eval (). cuda ()
# create example data
x = torch . ones (( 1 , 3 , 224 , 224 )). cuda ()
# convert to TensorRT feeding sample data as input
model_trt = torch2trt ( model , [ x ]) เราสามารถเรียกใช้ TRTModule ที่ส่งคืนได้เช่นเดียวกับรุ่น pytorch ดั้งเดิม
y = model ( x )
y_trt = model_trt ( x )
# check the output against PyTorch
print ( torch . max ( torch . abs ( y - y_trt ))) เราสามารถบันทึกโมเดลเป็น state_dict
torch . save ( model_trt . state_dict (), 'alexnet_trt.pth' ) เราสามารถโหลดโมเดลที่บันทึกไว้ใน TRTModule
from torch2trt import TRTModule
model_trt = TRTModule ()
model_trt . load_state_dict ( torch . load ( 'alexnet_trt.pth' ))เราทดสอบตัวแปลงกับโมเดลเหล่านี้โดยใช้สคริปต์ test.sh คุณสามารถสร้างผลลัพธ์ได้โดยการโทร
./test.sh TEST_OUTPUT.mdผลลัพธ์ด้านล่างแสดงปริมาณงานใน FPS คุณสามารถค้นหาเอาท์พุทดิบซึ่งรวมถึงเวลาแฝงในโฟลเดอร์มาตรฐาน
| แบบอย่าง | นาโน (Pytorch) | นาโน (Tensorrt) | Xavier (Pytorch) | Xavier (Tensorrt) |
|---|---|---|---|---|
| Alexnet | 46.4 | 69.9 | 250 | 580 |
| Squeezenet1_0 | 44 | 137 | 130 | 890 |
| Squeezenet1_1 | 76.6 | 248 | 132 | 1390 |
| resnet18 | 29.4 | 90.2 | 140 | 712 |
| resnet34 | 15.5 | 50.7 | 79.2 | 393 |
| resnet50 | 12.4 | 34.2 | 55.5 | 312 |
| resnet101 | 7.18 | 19.9 | 28.5 | 170 |
| Resnet152 | 4.96 | 14.1 | 18.9 | 121 |
| Densenet121 | 11.5 | 41.9 | 23.0 | 168 |
| Densenet169 | 8.25 | 33.2 | 16.3 | 118 |
| Densenet201 | 6.84 | 25.4 | 13.3 | 90.9 |
| Densenet161 | 4.71 | 15.6 | 17.2 | 82.4 |
| VGG11 | 8.9 | 18.3 | 85.2 | 201 |
| VGG13 | 6.53 | 14.7 | 71.9 | 166 |
| VGG16 | 5.09 | 11.9 | 61.7 | 139 |
| VGG19 | 54.1 | 121 | ||
| vgg11_bn | 8.74 | 18.4 | 81.8 | 201 |
| vgg13_bn | 6.31 | 14.8 | 68.0 | 166 |
| vgg16_bn | 4.96 | 12.0 | 58.5 | 140 |
| vgg19_bn | 51.4 | 121 |
หมายเหตุ: Torch2trt ขึ้นอยู่กับ Tensorrt Python API ใน Jetson สิ่งนี้รวมอยู่ใน Jetpack ล่าสุด สำหรับเดสก์ท็อปโปรดติดตามคู่มือการติดตั้ง Tensorrt คุณอาจลองติดตั้ง Torch2trt ภายในหนึ่งในคอนเทนเนอร์ NGC Pytorch Docker สำหรับเดสก์ท็อปหรือ Jetson
หากต้องการติดตั้งไลบรารี Torch2trt Python ให้โทรไปที่สิ่งต่อไปนี้
git clone https://github.com/NVIDIA-AI-IOT/torch2trt
cd torch2trt
python setup.py installหากต้องการติดตั้งไลบรารีปลั๊กอิน Torch2trt ให้โทรไปที่สิ่งต่อไปนี้
cmake -B build . && cmake --build build --target install && ldconfigซึ่งรวมถึงการสนับสนุนสำหรับเลเยอร์บางชั้นซึ่งอาจไม่ได้รับการสนับสนุนโดย Tensorrt เมื่อพบไลบรารีนี้ในระบบตัวแปลงเลเยอร์ที่เกี่ยวข้องใน Torch2trt จะถูกเปิดใช้งานโดยปริยาย
หมายเหตุ: Torch2trt ตอนนี้ดูแลปลั๊กอินเป็นห้องสมุดอิสระที่รวบรวมด้วย CMake สิ่งนี้ทำให้เครื่องยนต์ Tensorrt ที่รวบรวมได้พกพามากขึ้น หากจำเป็นปลั๊กอินที่เลิกใช้แล้ว (ซึ่งขึ้นอยู่กับ pytorch) อาจยังคงติดตั้งโดยการเรียกใช้
python setup.py install --plugins
ในการติดตั้ง Torch2trt กับชุมชนทดลองมีส่วนร่วมคุณสมบัติภายใต้ torch2trt.contrib เช่นการฝึกอบรมเชิงปริมาณ (QAT) ( requires TensorRT>=7.0 ) โทรไปที่สิ่งต่อไปนี้
git clone https://github.com/NVIDIA-AI-IOT/torch2trt
cd torch2trt/scripts
bash build_contrib.sh สิ่งนี้ช่วยให้คุณสามารถเรียกใช้ตัวอย่าง QAT ที่อยู่ที่นี่
ตัวแปลงนี้ใช้งานได้โดยการแนบฟังก์ชั่นการแปลง (เช่น convert_ReLU ) กับการเรียกใช้ฟังก์ชัน pytorch ดั้งเดิม (เช่น torch.nn.ReLU.forward ) ข้อมูลอินพุตตัวอย่างจะถูกส่งผ่านเครือข่ายเช่นเดียวกับก่อนหน้านี้ยกเว้นตอนนี้เมื่อใดก็ตามที่มีฟังก์ชั่นที่ลงทะเบียน ( torch.nn.ReLU.forward ) จะพบตัวแปลงที่สอดคล้องกัน ( convert_ReLU ) หลังจากนั้น ตัวแปลงถูกส่งผ่านข้อโต้แย้งและคำสั่งส่งคืนของฟังก์ชัน pytorch ดั้งเดิมเช่นเดียวกับเครือข่าย tensorrt ที่กำลังสร้าง เทนเซอร์อินพุตไปยังฟังก์ชั่น pytorch ดั้งเดิมได้รับการแก้ไขเพื่อให้มีแอตทริบิวต์ _trt ซึ่งเป็นคู่ tensorrt ไปยัง Pytorch Tensor ฟังก์ชั่นการแปลงใช้ _trt นี้เพื่อเพิ่มเลเยอร์ไปยังเครือข่าย Tensorrt จากนั้นตั้งค่าแอตทริบิวต์ _trt สำหรับเทนเซอร์เอาต์พุตที่เกี่ยวข้อง เมื่อโมเดลถูกเรียกใช้งานอย่างสมบูรณ์เทนเซอร์สุดท้ายจะถูกทำเครื่องหมายว่าเป็นเอาต์พุตของเครือข่าย Tensorrt และเครื่องยนต์ Tensorrt ที่ได้รับการปรับปรุงจะถูกสร้างขึ้น
ที่นี่เราแสดงวิธีเพิ่มตัวแปลงสำหรับโมดูล ReLU โดยใช้ Tensorrt Python API
import tensorrt as trt
from torch2trt import tensorrt_converter
@ tensorrt_converter ( 'torch.nn.ReLU.forward' )
def convert_ReLU ( ctx ):
input = ctx . method_args [ 1 ]
output = ctx . method_return
layer = ctx . network . add_activation ( input = input . _trt , type = trt . ActivationType . RELU )
output . _trt = layer . get_output ( 0 ) ตัวแปลงใช้อาร์กิวเมนต์หนึ่งรายการซึ่งจะมี ConversionContext ซึ่งจะมีสิ่งต่อไปนี้
ctx.network - เครือข่าย Tensorrt ที่กำลังสร้างขึ้น
ctx.method_args - อาร์กิวเมนต์ตำแหน่งที่ส่งผ่านไปยังฟังก์ชัน pytorch ที่ระบุ แอตทริบิวต์ _trt ถูกตั้งค่าสำหรับเทนเซอร์อินพุตที่เกี่ยวข้อง
ctx.method_kwargs - อาร์กิวเมนต์คำหลักที่ส่งผ่านไปยังฟังก์ชัน pytorch ที่ระบุ
ctx.method_return - ค่าที่ส่งคืนโดยฟังก์ชัน pytorch ที่ระบุ ตัวแปลงจะต้องตั้งค่าแอตทริบิวต์ _trt ที่เกี่ยวข้อง
โปรดดูโฟลเดอร์นี้สำหรับตัวอย่างเพิ่มเติม
JETBOT - หุ่นยนต์ AI เพื่อการศึกษาโดยใช้ Nvidia Jetson Nano
Jetracer - AI Racecar การศึกษาโดยใช้ Nvidia Jetson Nano
Jetcam - อินเทอร์เฟซกล้อง Python ที่ใช้งานง่ายสำหรับ Nvidia Jetson
JETCARD - ภาพการ์ด SD สำหรับการเขียนโปรแกรมเว็บโครงการ AI กับ Nvidia Jetson Nano
