- 1.python类的序列化(显然不太靠谱),我试了一下,果然不太靠谱
- 2.参考tensorRT官方文档
- 3.参考torch2trt官方git
# 序列化
serialized_engine = engine.serialize()
# 序列化并保存
with open(“sample.engine”, “wb”) as f:
f.write(engine.serialize())
# 反序列化
with trt.Runtime(TRT_LOGGER) as runtime:
engine = runtime.deserialize_cuda_engine(serialized_engine)
# 从文件中反序列化
with open(“sample.engine”, “rb”) as f, trt.Runtime(TRT_LOGGER) as runtime:
engine = runtime.deserialize_cuda_engine(f.read())
def torch2trt(module,
inputs,
input_names=None,
output_names=None,
log_level=trt.Logger.ERROR,
max_batch_size=1,
fp16_mode=False,
max_workspace_size=0,
strict_type_constraints=False,
keep_network=True,
int8_mode=False,
int8_calib_dataset=None,
int8_calib_algorithm=DEFAULT_CALIBRATION_ALGORITHM):
inputs_in = inputs
# copy inputs to avoid modifications to source data
inputs = [tensor.clone()[0:1] for tensor in inputs] # only run single entry
logger = trt.Logger(log_level)
builder = trt.Builder(logger)
network = builder.create_network()
with ConversionContext(network) as ctx:
if isinstance(inputs, list):
inputs = tuple(inputs)
if not isinstance(inputs, tuple):
inputs = (inputs, )
ctx.add_inputs(inputs, input_names)
outputs = module(*inputs)
if not isinstance(outputs, tuple) and not isinstance(outputs, list):
outputs = (outputs, )
ctx.mark_outputs(outputs, output_names)
builder.max_workspace_size = max_workspace_size
builder.fp16_mode = fp16_mode
builder.max_batch_size = max_batch_size
builder.strict_type_constraints = strict_type_constraints
if int8_mode:
# default to use input tensors for calibration
if int8_calib_dataset is None:
int8_calib_dataset = TensorBatchDataset(inputs_in)
builder.int8_mode = True
# @TODO(jwelsh): Should we set batch_size=max_batch_size? Need to investigate memory consumption
builder.int8_calibrator = DatasetCalibrator(inputs, int8_calib_dataset, batch_size=1, algorithm=int8_calib_algorithm)
engine = builder.build_cuda_engine(network)
module_trt = TRTModule(engine, ctx.input_names, ctx.output_names)
if keep_network:
module_trt.network = network
return module_trt
engine = builder.build_cuda_engine(network)
,简单来说,builder 负责构造网络,而engine也是由builder给build出来的serialized_engine = model_trt.engine.serialize()
with open("speed.engine", "wb") as f:
f.write(serialized_engine)
# 强行转换会报这个错误
Traceback (most recent call last):
File "/home/nano/Desktop/YOLOv3-Torch2TRT/mydetect.py", line 190, in <module>
detections = non_max_suppression(detections, opt.conf_thres, opt.nms_thres, method=2)
File "/home/nano/Desktop/YOLOv3-Torch2TRT/utils/utils.py", line 254, in non_max_suppression
image_pred = image_pred[(-score).argsort()]
IndexError: too many indices for tensor of dimension 2
# 序列化
torch.save(model_trt.state_dict(), 'alexnet_trt.pth')
# 读取
from torch2trt import TRTModule
model_trt = TRTModule()
model_trt.load_state_dict(torch.load('alexnet_trt.pth'))
序列化没测,反序列化测了一下
可以正常使用