Spaces:
Build error
Build error
File size: 4,279 Bytes
d7a991a |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 |
# Copyright (c) OpenMMLab. All rights reserved.
import torch
import torch.nn as nn
from mmpose.core import build_optimizers
class ExampleModel(nn.Module):
def __init__(self):
super().__init__()
self.model1 = nn.Conv2d(3, 8, kernel_size=3)
self.model2 = nn.Conv2d(3, 4, kernel_size=3)
def forward(self, x):
return x
def test_build_optimizers():
base_lr = 0.0001
base_wd = 0.0002
momentum = 0.9
# basic config with ExampleModel
optimizer_cfg = dict(
model1=dict(
type='SGD', lr=base_lr, weight_decay=base_wd, momentum=momentum),
model2=dict(
type='SGD', lr=base_lr, weight_decay=base_wd, momentum=momentum))
model = ExampleModel()
optimizers = build_optimizers(model, optimizer_cfg)
param_dict = dict(model.named_parameters())
assert isinstance(optimizers, dict)
for i in range(2):
optimizer = optimizers[f'model{i+1}']
param_groups = optimizer.param_groups[0]
assert isinstance(optimizer, torch.optim.SGD)
assert optimizer.defaults['lr'] == base_lr
assert optimizer.defaults['momentum'] == momentum
assert optimizer.defaults['weight_decay'] == base_wd
assert len(param_groups['params']) == 2
assert torch.equal(param_groups['params'][0],
param_dict[f'model{i+1}.weight'])
assert torch.equal(param_groups['params'][1],
param_dict[f'model{i+1}.bias'])
# basic config with Parallel model
model = torch.nn.DataParallel(ExampleModel())
optimizers = build_optimizers(model, optimizer_cfg)
param_dict = dict(model.named_parameters())
assert isinstance(optimizers, dict)
for i in range(2):
optimizer = optimizers[f'model{i+1}']
param_groups = optimizer.param_groups[0]
assert isinstance(optimizer, torch.optim.SGD)
assert optimizer.defaults['lr'] == base_lr
assert optimizer.defaults['momentum'] == momentum
assert optimizer.defaults['weight_decay'] == base_wd
assert len(param_groups['params']) == 2
assert torch.equal(param_groups['params'][0],
param_dict[f'module.model{i+1}.weight'])
assert torch.equal(param_groups['params'][1],
param_dict[f'module.model{i+1}.bias'])
# basic config with ExampleModel (one optimizer)
optimizer_cfg = dict(
type='SGD', lr=base_lr, weight_decay=base_wd, momentum=momentum)
model = ExampleModel()
optimizer = build_optimizers(model, optimizer_cfg)
param_dict = dict(model.named_parameters())
assert isinstance(optimizers, dict)
param_groups = optimizer.param_groups[0]
assert isinstance(optimizer, torch.optim.SGD)
assert optimizer.defaults['lr'] == base_lr
assert optimizer.defaults['momentum'] == momentum
assert optimizer.defaults['weight_decay'] == base_wd
assert len(param_groups['params']) == 4
assert torch.equal(param_groups['params'][0], param_dict['model1.weight'])
assert torch.equal(param_groups['params'][1], param_dict['model1.bias'])
assert torch.equal(param_groups['params'][2], param_dict['model2.weight'])
assert torch.equal(param_groups['params'][3], param_dict['model2.bias'])
# basic config with Parallel model (one optimizer)
model = torch.nn.DataParallel(ExampleModel())
optimizer = build_optimizers(model, optimizer_cfg)
param_dict = dict(model.named_parameters())
assert isinstance(optimizers, dict)
param_groups = optimizer.param_groups[0]
assert isinstance(optimizer, torch.optim.SGD)
assert optimizer.defaults['lr'] == base_lr
assert optimizer.defaults['momentum'] == momentum
assert optimizer.defaults['weight_decay'] == base_wd
assert len(param_groups['params']) == 4
assert torch.equal(param_groups['params'][0],
param_dict['module.model1.weight'])
assert torch.equal(param_groups['params'][1],
param_dict['module.model1.bias'])
assert torch.equal(param_groups['params'][2],
param_dict['module.model2.weight'])
assert torch.equal(param_groups['params'][3],
param_dict['module.model2.bias'])
|