github
/
ColossalAI
mirror of https://github.com/hpcaitech/ColossalAI
1
0
Fork 0
Code Issues Projects Releases Wiki Activity

[test] fixed codefactor format report (#4026)

develop
Frank Lee 2023-06-16 18:23:02 +08:00 committed by GitHub
parent 822c3d4d66
commit a5883aa790
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
1 changed files with 14 additions and 4 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

18
tests/test_checkpoint_io/test_general_checkpoint_io.py
View File

@ -165,8 +165,13 @@ def test_sharded_optimizer_multiple_param_groups():
# create a model and optimizer # create a model and optimizer
model = resnet18() model = resnet18()
optimizer = Adam([{'params': model.layer1.parameters()}, \ optimizer = Adam([{
{'params': model.layer2.parameters(), 'lr': 0.002}], lr=0.001) 'params': model.layer1.parameters()
}, {
'params': model.layer2.parameters(),
'lr': 0.002
}],
lr=0.001)
# create test data sample # create test data sample
x = torch.randn(1, 3, 224, 224) x = torch.randn(1, 3, 224, 224)
@ -189,8 +194,13 @@ def test_sharded_optimizer_multiple_param_groups():
# create new model # create new model
new_model = resnet18() new_model = resnet18()
new_optimizer = Adam([{'params': new_model.layer1.parameters()}, \ new_optimizer = Adam([{
{'params': new_model.layer2.parameters(), 'lr': 0.002}], lr=0.001) 'params': new_model.layer1.parameters()
}, {
'params': new_model.layer2.parameters(),
'lr': 0.002
}],
lr=0.001)
ckpt_io.load_model(new_model, str(model_ckpt_dir.name), strict=True) ckpt_io.load_model(new_model, str(model_ckpt_dir.name), strict=True)
ckpt_io.load_optimizer(new_optimizer, str(optimizer_ckpt_dir.name)) ckpt_io.load_optimizer(new_optimizer, str(optimizer_ckpt_dir.name))