This commit is contained in:
RunningLeon 2023-06-08 16:33:35 +08:00
parent 1d10af990d
commit 2273a00ff5
3 changed files with 5 additions and 6 deletions

View File

@ -34,7 +34,7 @@ def get_deploy_cfg():
backend_config=dict(
type='tensorrt',
common_config=dict(
fp16_mode=False, max_workspace_size=1 << 30),
fp16_mode=False, max_workspace_size=1 << 20),
model_inputs=[
dict(
input_shapes=dict(

View File

@ -403,7 +403,7 @@ def test_forward_of_fpnc(backend: Backend):
dict(
backend_config=dict(
type=backend.value,
common_config=dict(max_workspace_size=1 << 30),
common_config=dict(max_workspace_size=1 << 20),
model_inputs=[
dict(
input_shapes=dict(
@ -550,7 +550,7 @@ def test_mmdet_wrapper__forward(backend):
dict(
backend_config=dict(
type=backend.value,
common_config=dict(max_workspace_size=1 << 30)),
common_config=dict(max_workspace_size=1 << 20)),
onnx_config=dict(
input_shape=None,
input_names=['inputs'],

View File

@ -119,7 +119,7 @@ class TestTensorRTExporter:
backend_config=dict(
type='tensorrt',
common_config=dict(
fp16_mode=False, max_workspace_size=1 << 28),
fp16_mode=False, max_workspace_size=1 << 20),
model_inputs=[
dict(
input_shapes=dict(
@ -132,14 +132,13 @@ class TestTensorRTExporter:
])))
])))
onnx_model = onnx.load(onnx_file_path)
work_dir, filename = os.path.split(trt_file_path)
trt_apis.onnx2tensorrt(
work_dir,
filename,
0,
deploy_cfg=deploy_cfg,
onnx_model=onnx_model)
onnx_model=onnx_file_path)
if expected_result is None and not isinstance(
model, onnx.onnx_ml_pb2.ModelProto):
with torch.no_grad():