Skip to content
GitLab
Menu
Projects
Groups
Snippets
Loading...
Help
Help
Support
Community forum
Keyboard shortcuts
?
Submit feedback
Contribute to GitLab
Sign in / Register
Toggle navigation
Menu
Open sidebar
wangsen
paddle_dbnet
Commits
a2a12fe4
Commit
a2a12fe4
authored
Nov 25, 2021
by
LDOUBLEV
Browse files
fix TRT8 core bug
parent
fa1a1256
Changes
1
Hide whitespace changes
Inline
Side-by-side
Showing
1 changed file
with
5 additions
and
1 deletion
+5
-1
tools/infer/utility.py
tools/infer/utility.py
+5
-1
No files found.
tools/infer/utility.py
View file @
a2a12fe4
...
@@ -173,6 +173,7 @@ def create_predictor(args, mode, logger):
...
@@ -173,6 +173,7 @@ def create_predictor(args, mode, logger):
config
.
enable_use_gpu
(
args
.
gpu_mem
,
0
)
config
.
enable_use_gpu
(
args
.
gpu_mem
,
0
)
if
args
.
use_tensorrt
:
if
args
.
use_tensorrt
:
config
.
enable_tensorrt_engine
(
config
.
enable_tensorrt_engine
(
workspace_size
=
1
<<
30
,
precision_mode
=
precision
,
precision_mode
=
precision
,
max_batch_size
=
args
.
max_batch_size
,
max_batch_size
=
args
.
max_batch_size
,
min_subgraph_size
=
args
.
min_subgraph_size
)
min_subgraph_size
=
args
.
min_subgraph_size
)
...
@@ -294,7 +295,10 @@ def create_predictor(args, mode, logger):
...
@@ -294,7 +295,10 @@ def create_predictor(args, mode, logger):
def
get_infer_gpuid
():
def
get_infer_gpuid
():
cmd
=
"nvidia-smi"
cmd
=
"nvidia-smi"
res
=
os
.
popen
(
cmd
).
readlines
()
try
:
res
=
os
.
popen
(
cmd
).
readlines
()
except
:
res
=
None
if
len
(
res
)
==
0
:
if
len
(
res
)
==
0
:
return
None
return
None
cmd
=
"env | grep CUDA_VISIBLE_DEVICES"
cmd
=
"env | grep CUDA_VISIBLE_DEVICES"
...
...
Write
Preview
Markdown
is supported
0%
Try again
or
attach a new file
.
Attach a file
Cancel
You are about to add
0
people
to the discussion. Proceed with caution.
Finish editing this message first!
Cancel
Please
register
or
sign in
to comment