YoloV4 keras 到 TensorRT

Question

我正在根据此存储库中的 keras 自定义训练 YoloV4： https://github.com/ taipingeric/yolo-v4-tf.keras

model = Yolov4(weight_path=weights,
               class_name_path=class_name_path)

model.load_model('saved_model/with es')
#model.predict('input.jpg')

conversion_params = trt.DEFAULT_TRT_CONVERSION_PARAMS._replace(
    precision_mode=trt.TrtPrecisionMode.FP16,
    max_workspace_size_bytes=4000000000,
    max_batch_size=4)

converter = trt.TrtGraphConverterV2(
    input_saved_model_dir=saved_model, conversion_params=conversion_params)
converter.convert()
converter.save(output_saved_model_dir)

然后使用 NVIDIA 的 TF20-TF-TRT 指南加载相同的模型：

saved_model_loaded = tf.saved_model.load(input_saved_model, tags=[tag_constants.SERVING])
signature_keys = list(saved_model_loaded.signatures.keys())
print(signature_keys)
infer = saved_model_loaded.signatures['serving_default']
print(infer.structured_outputs)

但是当我尝试使用它来推断时，

labeling = infer(x)

我得到了以下错误：

2022-03-02 10:52:18.993365: W tensorflow/compiler/tf2tensorrt/kernels/trt_engine_op.cc:628] TF-TRT Warning: Engine retrieval for input shapes: [[1,416,416,3]] failed. Running native segment for PartitionedCall/TRTEngineOp_0_1
2022-03-02 10:52:19.035613: E tensorflow/stream_executor/cuda/cuda_blas.cc:226] failed to create cublas handle: CUBLAS_STATUS_ALLOC_FAILED
2022-03-02 10:52:19.035702: W tensorflow/core/framework/op_kernel.cc:1763] OP_REQUIRES failed at conv_ops.cc:1106 : Not found: No algorithm worked!
2022-03-02 10:52:19.035770: W tensorflow/core/framework/op_kernel.cc:1763] OP_REQUIRES failed at trt_engine_op.cc:400 : Not found: No algorithm worked!
         [[{{node StatefulPartitionedCall/model_1/conv2d/Conv2D}}]]
Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
  File "/mnt/d/Testing/research/yolo-v4/yolo-v4/lib/python3.6/site-packages/tensorflow/python/eager/function.py", line 1669, in __call__
    return self._call_impl(args, kwargs)
  File "/mnt/d/Testing/research/yolo-v4/yolo-v4/lib/python3.6/site-packages/tensorflow/python/eager/function.py", line 1679, in _call_impl
    cancellation_manager)
  File "/mnt/d/Testing/research/yolo-v4/yolo-v4/lib/python3.6/site-packages/tensorflow/python/eager/function.py", line 1762, in _call_with_structured_signature
    cancellation_manager=cancellation_manager)
  File "/mnt/d/Testing/research/yolo-v4/yolo-v4/lib/python3.6/site-packages/tensorflow/python/saved_model/load.py", line 116, in _call_flat
    cancellation_manager)
  File "/mnt/d/Testing/research/yolo-v4/yolo-v4/lib/python3.6/site-packages/tensorflow/python/eager/function.py", line 1919, in _call_flat
    ctx, args, cancellation_manager=cancellation_manager))
  File "/mnt/d/Testing/research/yolo-v4/yolo-v4/lib/python3.6/site-packages/tensorflow/python/eager/function.py", line 560, in call
    ctx=ctx)
  File "/mnt/d/Testing/research/yolo-v4/yolo-v4/lib/python3.6/site-packages/tensorflow/python/eager/execute.py", line 60, in quick_execute
    inputs, attrs, num_outputs)
tensorflow.python.framework.errors_impl.NotFoundError:  No algorithm worked!
         [[{{node StatefulPartitionedCall/model_1/conv2d/Conv2D}}]]
         [[PartitionedCall/TRTEngineOp_0_1]] [Op:__inference_signature_wrapper_130807]

Function call stack:
signature_wrapper