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
ColossalAI
Commits
9e512938
Commit
9e512938
authored
Jul 19, 2023
by
Zirui Zhu
Committed by
binmakeswell
Jul 26, 2023
Browse files
[NFC] polish applications/Chat/coati/trainer/strategies/base.py code style (#4278)
parent
c972d653
Changes
1
Hide whitespace changes
Inline
Side-by-side
Showing
1 changed file
with
4 additions
and
18 deletions
+4
-18
applications/Chat/coati/trainer/strategies/base.py
applications/Chat/coati/trainer/strategies/base.py
+4
-18
No files found.
applications/Chat/coati/trainer/strategies/base.py
View file @
9e512938
...
@@ -79,8 +79,7 @@ class Strategy(ABC):
...
@@ -79,8 +79,7 @@ class Strategy(ABC):
model
,
optimizer
=
arg
model
,
optimizer
=
arg
except
ValueError
:
except
ValueError
:
raise
RuntimeError
(
f
'Expect (model, optimizer) pair, got a tuple with size "
{
len
(
arg
)
}
"'
)
raise
RuntimeError
(
f
'Expect (model, optimizer) pair, got a tuple with size "
{
len
(
arg
)
}
"'
)
model
,
optimizer
,
*
_
=
self
.
booster
.
boost
(
model
=
model
,
model
,
optimizer
,
*
_
=
self
.
booster
.
boost
(
model
=
model
,
optimizer
=
optimizer
)
optimizer
=
optimizer
)
rets
.
append
((
model
,
optimizer
))
rets
.
append
((
model
,
optimizer
))
elif
isinstance
(
arg
,
Dict
):
elif
isinstance
(
arg
,
Dict
):
model
,
optimizer
,
criterion
,
dataloader
,
lr_scheduler
=
self
.
booster
.
boost
(
**
arg
)
model
,
optimizer
,
criterion
,
dataloader
,
lr_scheduler
=
self
.
booster
.
boost
(
**
arg
)
...
@@ -90,10 +89,7 @@ class Strategy(ABC):
...
@@ -90,10 +89,7 @@ class Strategy(ABC):
dataloader
=
dataloader
,
dataloader
=
dataloader
,
lr_scheduler
=
lr_scheduler
)
lr_scheduler
=
lr_scheduler
)
# remove None values
# remove None values
boost_result
=
{
boost_result
=
{
key
:
value
for
key
,
value
in
boost_result
.
items
()
if
value
is
not
None
}
key
:
value
for
key
,
value
in
boost_result
.
items
()
if
value
is
not
None
}
rets
.
append
(
boost_result
)
rets
.
append
(
boost_result
)
else
:
else
:
raise
RuntimeError
(
f
'Type
{
type
(
arg
)
}
is not supported'
)
raise
RuntimeError
(
f
'Type
{
type
(
arg
)
}
is not supported'
)
...
@@ -112,23 +108,13 @@ class Strategy(ABC):
...
@@ -112,23 +108,13 @@ class Strategy(ABC):
"""
"""
return
model
return
model
def
save_model
(
self
,
def
save_model
(
self
,
model
:
nn
.
Module
,
path
:
str
,
only_rank0
:
bool
=
True
,
**
kwargs
)
->
None
:
model
:
nn
.
Module
,
path
:
str
,
only_rank0
:
bool
=
True
,
**
kwargs
)
->
None
:
self
.
booster
.
save_model
(
model
,
path
,
shard
=
not
only_rank0
,
**
kwargs
)
self
.
booster
.
save_model
(
model
,
path
,
shard
=
not
only_rank0
,
**
kwargs
)
def
load_model
(
self
,
model
:
nn
.
Module
,
path
:
str
,
strict
:
bool
=
True
)
->
None
:
def
load_model
(
self
,
model
:
nn
.
Module
,
path
:
str
,
strict
:
bool
=
True
)
->
None
:
self
.
booster
.
load_model
(
model
,
path
,
strict
)
self
.
booster
.
load_model
(
model
,
path
,
strict
)
def
save_optimizer
(
self
,
def
save_optimizer
(
self
,
optimizer
:
Optimizer
,
path
:
str
,
only_rank0
:
bool
=
False
,
**
kwargs
)
->
None
:
optimizer
:
Optimizer
,
path
:
str
,
only_rank0
:
bool
=
False
,
**
kwargs
)
->
None
:
self
.
booster
.
save_optimizer
(
optimizer
,
path
,
shard
=
not
only_rank0
,
**
kwargs
)
self
.
booster
.
save_optimizer
(
optimizer
,
path
,
shard
=
not
only_rank0
,
**
kwargs
)
def
load_optimizer
(
self
,
optimizer
:
Optimizer
,
path
:
str
)
->
None
:
def
load_optimizer
(
self
,
optimizer
:
Optimizer
,
path
:
str
)
->
None
:
...
...
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