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
MinerU
Commits
84b3c3bb
Commit
84b3c3bb
authored
Jul 10, 2024
by
zhaoxiaomeng
Browse files
update: add mfr cost time each batch of dataloader
parent
9302caa3
Changes
1
Hide whitespace changes
Inline
Side-by-side
Showing
1 changed file
with
3 additions
and
0 deletions
+3
-0
magic_pdf/model/pdf_extract_kit.py
magic_pdf/model/pdf_extract_kit.py
+3
-0
No files found.
magic_pdf/model/pdf_extract_kit.py
View file @
84b3c3bb
...
...
@@ -144,9 +144,12 @@ class CustomPEKModel:
dataloader
=
DataLoader
(
dataset
,
batch_size
=
128
,
num_workers
=
0
)
mfr_res
=
[]
for
imgs
in
dataloader
:
start
=
time
.
time
()
imgs
=
imgs
.
to
(
self
.
device
)
output
=
self
.
mfr_model
.
generate
({
'image'
:
imgs
})
mfr_res
.
extend
(
output
[
'pred_str'
])
cost
=
time
.
time
()
-
start
logger
.
info
(
f
"batch size:
{
len
(
imgs
)
}
, cost time:
{
round
(
cost
,
2
)
}
"
)
for
res
,
latex
in
zip
(
latex_filling_list
,
mfr_res
):
res
[
'latex'
]
=
latex_rm_whitespace
(
latex
)
b
=
time
.
time
()
...
...
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