Как преобразовать YOLOv8 в TensorRT с помощью torch2trt? ⇐ Python

[Anonymous]

Как видно из заголовка, в следующем коде есть ошибки или упущения, что приводит к неудачному преобразованию.

Код: Выделить всё

import torch
from torch2trt import torch2trt
from ultralytics import YOLO

if __name__ == '__main__':
model = YOLO('yolov8s_640.pt')

x = torch.ones((1, 3, 640, 640))
x = x.cuda()

# # convert to TensorRT feeding sample data as input
model_trt = torch2trt(model, [x])

torch.save(model_trt.state_dict(), './yolov8s_640_trt.pth')

По приведенному ниже методу кажется, что требуется загрузка самой модели, так как файл yolov8s_160.pt содержит только веса. Как загрузить модель yolov8?

Код: Выделить всё

import torch
from torch2trt import torch2trt
from ultralytics import YOLO

if __name__ == '__main__':
model = torch.load('yolov8s_640.pt')
model = model.eval().cuda()

x = torch.ones((1, 3, 640, 640))
x = x.cuda()

# # convert to TensorRT feeding sample data as input
model_trt = torch2trt(model, [x])

torch.save(model_trt.state_dict(), './yolov8s_640_trt.pth')

Может ли кто-нибудь из опытных предложить правильный подход?

Подробнее здесь: https://stackoverflow.com/questions/784 ... -torch2trt