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
chenpangpang
transformers
Commits
11792d78
Unverified
Commit
11792d78
authored
Jul 27, 2020
by
Sam Shleifer
Committed by
GitHub
Jul 27, 2020
Browse files
CL util to convert models to fp16 before upload (#5953)
parent
4302ace5
Changes
1
Hide whitespace changes
Inline
Side-by-side
Showing
1 changed file
with
21 additions
and
0 deletions
+21
-0
examples/seq2seq/convert_model_to_fp16.py
examples/seq2seq/convert_model_to_fp16.py
+21
-0
No files found.
examples/seq2seq/convert_model_to_fp16.py
0 → 100644
View file @
11792d78
from
typing
import
Union
import
fire
import
torch
from
tqdm
import
tqdm
def
convert
(
src_path
:
str
,
map_location
:
str
=
"cpu"
,
save_path
:
Union
[
str
,
None
]
=
None
)
->
None
:
"""Convert a pytorch_model.bin or model.pt file to torch.float16 for faster downloads, less disk space."""
state_dict
=
torch
.
load
(
src_path
,
map_location
=
map_location
)
for
k
,
v
in
tqdm
(
state_dict
.
items
()):
if
not
isinstance
(
v
,
torch
.
Tensor
):
raise
TypeError
(
"FP16 conversion only works on paths that are saved state dics, like pytorch_model.bin"
)
state_dict
[
k
]
=
v
.
half
()
if
save_path
is
None
:
# overwrite src_path
save_path
=
src_path
torch
.
save
(
state_dict
,
save_path
)
if
__name__
==
"__main__"
:
fire
.
Fire
(
convert
)
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