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
OpenDAS
MMCV
Commits
f48241a6
Unverified
Commit
f48241a6
authored
Jun 17, 2020
by
Wang Xinjiang
Committed by
GitHub
Jun 17, 2020
Browse files
load_checkpoint support normal dict checkpoints (#351)
* load_checkpoint support normal dict checkpoints * comments
parent
97f9efd8
Changes
1
Hide whitespace changes
Inline
Side-by-side
Showing
1 changed file
with
6 additions
and
5 deletions
+6
-5
mmcv/runner/checkpoint.py
mmcv/runner/checkpoint.py
+6
-5
No files found.
mmcv/runner/checkpoint.py
View file @
f48241a6
...
@@ -222,14 +222,15 @@ def load_checkpoint(model,
...
@@ -222,14 +222,15 @@ def load_checkpoint(model,
dict or OrderedDict: The loaded checkpoint.
dict or OrderedDict: The loaded checkpoint.
"""
"""
checkpoint
=
_load_checkpoint
(
filename
,
map_location
)
checkpoint
=
_load_checkpoint
(
filename
,
map_location
)
# OrderedDict is a subclass of dict
if
not
isinstance
(
checkpoint
,
dict
):
raise
RuntimeError
(
f
'No state_dict found in checkpoint file
{
filename
}
'
)
# get state_dict from checkpoint
# get state_dict from checkpoint
if
isinstance
(
checkpoint
,
OrderedDict
):
if
'state_dict'
in
checkpoint
:
state_dict
=
checkpoint
elif
isinstance
(
checkpoint
,
dict
)
and
'state_dict'
in
checkpoint
:
state_dict
=
checkpoint
[
'state_dict'
]
state_dict
=
checkpoint
[
'state_dict'
]
else
:
else
:
raise
RuntimeError
(
state_dict
=
checkpoint
f
'No state_dict found in checkpoint file
{
filename
}
'
)
# strip prefix of state_dict
# strip prefix of state_dict
if
list
(
state_dict
.
keys
())[
0
].
startswith
(
'module.'
):
if
list
(
state_dict
.
keys
())[
0
].
startswith
(
'module.'
):
state_dict
=
{
k
[
7
:]:
v
for
k
,
v
in
checkpoint
[
'state_dict'
].
items
()}
state_dict
=
{
k
[
7
:]:
v
for
k
,
v
in
checkpoint
[
'state_dict'
].
items
()}
...
...
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