当前位置：首页 > news >正文

Ultralytics导出的Engine模型直接加载报错

news 来源：原创 2025/5/19 0:32:46

方法1: 注销掉exporter.py文件夹对应的部分

with build(network, config) as engine, open(f, "wb") as t:
    # Metadata
    meta = json.dumps(self.metadata)
    t.write(len(meta).to_bytes(4, byteorder="little", signed=True))
    t.write(meta.encode())
    # Model
    t.write(engine if is_trt10 else engine.serialize())

替换为

with build(network, config) as engine, open(f, "wb") as t:
    # Metadata
    # meta = json.dumps(self.metadata)
    # t.write(len(meta).to_bytes(4, byteorder="little", signed=True))
    # t.write(meta.encode())
    # Model
    t.write(engine if is_trt10 else engine.serialize())

方法2：加载模型时去掉元数据
python环境下

with open(engine_path, 'rb') as f, trt.Runtime(TRT_LOGGER) as runtime:
    engine = runtime.deserialize_cuda_engine(f.read())

替换为

with open(engine_path, 'rb') as f:
    metadata_length_bytes = f.read(4)
    metadata_length = int.from_bytes(metadata_length_bytes, byteorder='little', signed=True)
    f.seek(metadata_length, 1)
    engine_data = f.read()
with open(engine_path, 'rb') as f, trt.Runtime(TRT_LOGGER) as runtime:
    engine = runtime.deserialize_cuda_engine(engine_data)

cpp环境下

std::ifstream file_ptr(model_path_engine, std::ios::binary);
size_t size = 0;
file_ptr.seekg(0, file_ptr.end);
size = file_ptr.tellg();
file_ptr.seekg(0, file_ptr.beg);
char* model_stream = new char[size];
file_ptr.read(model_stream, size);
file_ptr.close();
nvinfer1::IRuntime* runtime = nvinfer1::createInferRuntime(logger);
nvinfer1::ICudaEngine* engine = runtime->deserializeCudaEngine(model_stream, size);
delete[] model_stream;

替换为

std::ifstream engineFile(model_path_engine, std::ios::binary);
int32_t metadataLength = 0;
engineFile.read(reinterpret_cast<char*>(&metadataLength), sizeof(metadataLength));
engineFile.seekg(metadataLength, std::ios::cur);
std::vector<char> engineData((std::istreambuf_iterator<char>(engineFile)),
	std::istreambuf_iterator<char>());
engineFile.close();
nvinfer1::IRuntime* runtime = nvinfer1::createInferRuntime(logger);
nvinfer1::ICudaEngine* engine = runtime->deserializeCudaEngine(engineData.data(), engineData.size(), nullptr);