Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

[Refactor] Change the meaning of momentum in EMA #1581

Merged
merged 13 commits into from
Jan 18, 2023
Merged

[Refactor] Change the meaning of momentum in EMA #1581

Show file tree
Hide file tree
Changes from 3 commits
Commits
Show all changes
13 commits
Select commit Hold shift + click to select a range
a7c3303
fix ema momentum meaning, results in codes and config changes
plyfager Jan 6, 2023
d76cebf
fix comments
plyfager Jan 9, 2023
ce6f364
Merge branch 'dev-1.x' of github.com:open-mmlab/mmediting into plyfag…
plyfager Jan 9, 2023
7477d49
Update mmedit/engine/hooks/ema.py
plyfager Jan 9, 2023
866d9d3
complete warning for old user
plyfager Jan 11, 2023
fadcef3
fix lint
plyfager Jan 11, 2023
d51abbf
Merge branch 'dev-1.x' into plyfager/ema-check
zengyh1900 Jan 16, 2023
f4694cd
fix ut
plyfager Jan 17, 2023
5d3d841
Merge branch 'dev-1.x' of github.com:open-mmlab/mmediting into plyfag…
plyfager Jan 17, 2023
8d7795a
Merge branch 'plyfager/ema-check' of github.com:plyfager/mmediting in…
plyfager Jan 17, 2023
450fa92
test warning
plyfager Jan 17, 2023
22a640e
pytest capture warning
plyfager Jan 18, 2023
d7fe8c0
fix lint
plyfager Jan 18, 2023
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 1 addition & 1 deletion configs/biggan/biggan-deep_cvt-hugging-face-rgb_imagenet1k-128x128.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -6,7 +6,7 @@
ema_config = dict(
type='ExponentialMovingAverage',
interval=1,
momentum=0.9999,
momentum=0.0001,
update_buffers=True,
start_iter=20000)

Expand Down
2 changes: 1 addition & 1 deletion configs/biggan/biggan-deep_cvt-hugging-face_rgb_imagenet1k-256x256.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -11,7 +11,7 @@
ema_config = dict(
type='ExponentialMovingAverage',
interval=1,
momentum=0.9999,
momentum=0.0001,
update_buffers=True,
start_iter=20000)

Expand Down
2 changes: 1 addition & 1 deletion configs/biggan/biggan-deep_cvt-hugging-face_rgb_imagenet1k-512x512.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -11,7 +11,7 @@
ema_config = dict(
type='ExponentialMovingAverage',
interval=1,
momentum=0.9999,
momentum=0.0001,
update_buffers=True,
start_iter=20000)

Expand Down
2 changes: 1 addition & 1 deletion configs/biggan/biggan_2xb25-500kiters_cifar10-32x32.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -7,7 +7,7 @@
ema_config = dict(
type='ExponentialMovingAverage',
interval=1,
momentum=0.9999,
momentum=0.0001,
start_iter=1000)

model = dict(
Expand Down
2 changes: 1 addition & 1 deletion configs/biggan/biggan_ajbrock-sn_8xb32-1500kiters_imagenet1k-128x128.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -8,7 +8,7 @@
ema_config = dict(
type='ExponentialMovingAverage',
interval=1,
momentum=0.9999,
momentum=0.0001,
update_buffers=True,
start_iter=20000)

Expand Down
2 changes: 1 addition & 1 deletion configs/biggan/biggan_cvt-BigGAN-PyTorch-rgb_imagenet1k-128x128.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -6,7 +6,7 @@
ema_config = dict(
type='ExponentialMovingAverage',
interval=1,
momentum=0.9999,
momentum=0.0001,
update_buffers=True,
start_iter=20000)

Expand Down
2 changes: 1 addition & 1 deletion ...s/positional_encoding_in_gans/mspie-stylegan2-config-c_c2_8xb3-1100kiters_ffhq-256-512.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -7,7 +7,7 @@
ema_config = dict(
type='ExponentialMovingAverage',
interval=1,
momentum=0.5**(32. / (ema_half_life * 1000.)))
momentum=1. - (0.5**(32. / (ema_half_life * 1000.))))

model = dict(
type='MSPIEStyleGAN2',
Expand Down
2 changes: 1 addition & 1 deletion ...s/positional_encoding_in_gans/mspie-stylegan2-config-d_c2_8xb3-1100kiters_ffhq-256-512.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -7,7 +7,7 @@
ema_config = dict(
type='ExponentialMovingAverage',
interval=1,
momentum=0.5**(32. / (ema_half_life * 1000.)))
momentum=1. - (0.5**(32. / (ema_half_life * 1000.))))

model = dict(
type='MSPIEStyleGAN2',
Expand Down
2 changes: 1 addition & 1 deletion ...s/positional_encoding_in_gans/mspie-stylegan2-config-e_c2_8xb3-1100kiters_ffhq-256-512.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -7,7 +7,7 @@
ema_config = dict(
type='ExponentialMovingAverage',
interval=1,
momentum=0.5**(32. / (ema_half_life * 1000.)))
momentum=1. - (0.5**(32. / (ema_half_life * 1000.))))

model = dict(
type='MSPIEStyleGAN2',
Expand Down
2 changes: 1 addition & 1 deletion .../positional_encoding_in_gans/mspie-stylegan2-config-f_c1_8xb2-1600kiters_ffhq-256-1024.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -29,7 +29,7 @@
ema_config = dict(
type='ExponentialMovingAverage',
interval=1,
momentum=0.5**(32. / (ema_half_life * 1000.)))
momentum=1. - (0.5**(32. / (ema_half_life * 1000.))))

model = dict(
type='MSPIEStyleGAN2',
Expand Down
2 changes: 1 addition & 1 deletion ...s/positional_encoding_in_gans/mspie-stylegan2-config-f_c2_8xb3-1100kiters_ffhq-256-512.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -7,7 +7,7 @@
ema_config = dict(
type='ExponentialMovingAverage',
interval=1,
momentum=0.5**(32. / (ema_half_life * 1000.)))
momentum=1. - (0.5**(32. / (ema_half_life * 1000.))))

model = dict(
type='MSPIEStyleGAN2',
Expand Down
2 changes: 1 addition & 1 deletion ...s/positional_encoding_in_gans/mspie-stylegan2-config-f_c2_8xb3-1100kiters_ffhq-256-896.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -7,7 +7,7 @@
ema_config = dict(
type='ExponentialMovingAverage',
interval=1,
momentum=0.5**(32. / (ema_half_life * 1000.)))
momentum=1. - (0.5**(32. / (ema_half_life * 1000.))))

model = dict(
type='MSPIEStyleGAN2',
Expand Down
2 changes: 1 addition & 1 deletion ...s/positional_encoding_in_gans/mspie-stylegan2-config-g_c1_8xb3-1100kiters_ffhq-256-512.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -7,7 +7,7 @@
ema_config = dict(
type='ExponentialMovingAverage',
interval=1,
momentum=0.5**(32. / (ema_half_life * 1000.)))
momentum=1. - (0.5**(32. / (ema_half_life * 1000.))))

model = dict(
type='MSPIEStyleGAN2',
Expand Down
2 changes: 1 addition & 1 deletion ...s/positional_encoding_in_gans/mspie-stylegan2-config-h_c2_8xb3-1100kiters_ffhq-256-512.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -7,7 +7,7 @@
ema_config = dict(
type='ExponentialMovingAverage',
interval=1,
momentum=0.5**(32. / (ema_half_life * 1000.)))
momentum=1. - (0.5**(32. / (ema_half_life * 1000.))))

model = dict(
type='MSPIEStyleGAN2',
Expand Down
2 changes: 1 addition & 1 deletion ...s/positional_encoding_in_gans/mspie-stylegan2-config-i_c2_8xb3-1100kiters_ffhq-256-512.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -7,7 +7,7 @@
ema_config = dict(
type='ExponentialMovingAverage',
interval=1,
momentum=0.5**(32. / (ema_half_life * 1000.)))
momentum=1. - (0.5**(32. / (ema_half_life * 1000.))))

model = dict(
type='MSPIEStyleGAN2',
Expand Down
2 changes: 1 addition & 1 deletion ...s/positional_encoding_in_gans/mspie-stylegan2-config-j_c2_8xb3-1100kiters_ffhq-256-512.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -7,7 +7,7 @@
ema_config = dict(
type='ExponentialMovingAverage',
interval=1,
momentum=0.5**(32. / (ema_half_life * 1000.)))
momentum=1. - (0.5**(32. / (ema_half_life * 1000.))))

model = dict(
type='MSPIEStyleGAN2',
Expand Down
2 changes: 1 addition & 1 deletion ...s/positional_encoding_in_gans/mspie-stylegan2-config-k_c2_8xb3-1100kiters_ffhq-256-512.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -7,7 +7,7 @@
ema_config = dict(
type='ExponentialMovingAverage',
interval=1,
momentum=0.5**(32. / (ema_half_life * 1000.)))
momentum=1. - (0.5**(32. / (ema_half_life * 1000.))))

model = dict(
type='MSPIEStyleGAN2',
Expand Down
2 changes: 1 addition & 1 deletion configs/positional_encoding_in_gans/stylegan2_c2_8xb3-1100kiters_ffhq-256x256.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -20,7 +20,7 @@
ema_config=dict(
type='ExponentialMovingAverage',
interval=1,
momentum=0.5**(32. / (ema_half_life * 1000.))))
momentum=1. - (0.5**(32. / (ema_half_life * 1000.)))))

optim_wrapper = dict(
generator=dict(
Expand Down
2 changes: 1 addition & 1 deletion configs/positional_encoding_in_gans/stylegan2_c2_8xb3-1100kiters_ffhq-512x512.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -20,7 +20,7 @@
ema_config=dict(
type='ExponentialMovingAverage',
interval=1,
momentum=0.5**(32. / (ema_half_life * 1000.))))
momentum=1. - (0.5**(32. / (ema_half_life * 1000.)))))

optim_wrapper = dict(
generator=dict(
Expand Down
2 changes: 1 addition & 1 deletion configs/real_basicvsr/realbasicvsr_wogan-c64b20-2x30x8_8xb2-lr1e-4-300k_reds.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -301,7 +301,7 @@
type='ExponentialMovingAverageHook',
module_keys=('generator_ema'),
interval=1,
interp_cfg=dict(momentum=0.999),
interp_cfg=dict(momentum=0.001),
)
]

Expand Down
2 changes: 1 addition & 1 deletion configs/real_esrgan/realesrnet_c64b23g32_4xb12-lr2e-4-1000k_df2k-ost.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -259,6 +259,6 @@
type='ExponentialMovingAverageHook',
module_keys=('generator_ema'),
interval=1,
interp_cfg=dict(momentum=0.999),
interp_cfg=dict(momentum=0.001),
)
]
3 changes: 2 additions & 1 deletion configs/styleganv1/styleganv1_ffhq-1024x1024_8xb4-25Mimgs.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -7,7 +7,8 @@
# MODEL
model_wrapper_cfg = dict(find_unused_parameters=True)
ema_half_life = 10. # G_smoothing_kimg
ema_config = dict(interval=1, momentum=0.5**(32. / (ema_half_life * 1000.)))
ema_config = dict(
interval=1, momentum=1. - (0.5**(32. / (ema_half_life * 1000.))))
model = dict(
generator=dict(out_size=1024),
discriminator=dict(in_size=1024),
Expand Down
3 changes: 2 additions & 1 deletion configs/styleganv1/styleganv1_ffhq-256x256_8xb4-25Mimgs.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -7,7 +7,8 @@
# MODEL
model_wrapper_cfg = dict(find_unused_parameters=True)
ema_half_life = 10. # G_smoothing_kimg
ema_config = dict(interval=1, momentum=0.5**(32. / (ema_half_life * 1000.)))
ema_config = dict(
interval=1, momentum=1. - (0.5**(32. / (ema_half_life * 1000.))))
model = dict(
generator=dict(out_size=256),
discriminator=dict(in_size=256),
Expand Down
2 changes: 1 addition & 1 deletion configs/styleganv2/stylegan2_c2_8xb4-800kiters_ffhq-256x256.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -20,7 +20,7 @@
ema_config=dict(
type='ExponentialMovingAverage',
interval=1,
momentum=0.5**(32. / (ema_half_life * 1000.))),
momentum=1. - (0.5**(32. / (ema_half_life * 1000.)))),
loss_config=dict(
r1_loss_weight=10. / 2. * d_reg_interval,
r1_interval=d_reg_interval,
Expand Down
2 changes: 1 addition & 1 deletion configs/styleganv2/stylegan2_c2_8xb4-800kiters_lsun-cat-256x256.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -20,7 +20,7 @@
ema_config=dict(
type='ExponentialMovingAverage',
interval=1,
momentum=0.5**(32. / (ema_half_life * 1000.))),
momentum=1. - (0.5**(32. / (ema_half_life * 1000.)))),
loss_config=dict(
r1_loss_weight=10. / 2. * d_reg_interval,
r1_interval=d_reg_interval,
Expand Down
2 changes: 1 addition & 1 deletion configs/styleganv2/stylegan2_c2_8xb4-800kiters_lsun-church-256x256.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -20,7 +20,7 @@
ema_config=dict(
type='ExponentialMovingAverage',
interval=1,
momentum=0.5**(32. / (ema_half_life * 1000.))),
momentum=1. - (0.5**(32. / (ema_half_life * 1000.)))),
loss_config=dict(
r1_loss_weight=10. / 2. * d_reg_interval,
r1_interval=d_reg_interval,
Expand Down
2 changes: 1 addition & 1 deletion configs/styleganv2/stylegan2_c2_8xb4-800kiters_lsun-horse-256x256.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -20,7 +20,7 @@
ema_config=dict(
type='ExponentialMovingAverage',
interval=1,
momentum=0.5**(32. / (ema_half_life * 1000.))),
momentum=1. - (0.5**(32. / (ema_half_life * 1000.)))),
loss_config=dict(
r1_loss_weight=10. / 2. * d_reg_interval,
r1_interval=d_reg_interval,
Expand Down
2 changes: 1 addition & 1 deletion configs/styleganv2/stylegan2_c2_8xb4_ffhq-1024x1024.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -20,7 +20,7 @@
ema_config=dict(
type='ExponentialMovingAverage',
interval=1,
momentum=0.5**(32. / (ema_half_life * 1000.))),
momentum=1. - (0.5**(32. / (ema_half_life * 1000.)))),
loss_config=dict(
r1_loss_weight=10. / 2. * d_reg_interval,
r1_interval=d_reg_interval,
Expand Down
2 changes: 1 addition & 1 deletion configs/styleganv2/stylegan2_c2_8xb4_lsun-car-384x512.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -18,7 +18,7 @@
ema_config=dict(
type='ExponentialMovingAverage',
interval=1,
momentum=0.5**(32. / (ema_half_life * 1000.))),
momentum=1. - (0.5**(32. / (ema_half_life * 1000.)))),
loss_config=dict(
r1_loss_weight=10. / 2. * d_reg_interval,
r1_interval=d_reg_interval,
Expand Down
29 changes: 18 additions & 11 deletions mmedit/engine/hooks/ema.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -69,22 +69,29 @@ def __init__(self,
getattr(self, interp_mode), **self.interp_cfg)

@staticmethod
def lerp(a, b, momentum=0.999, momentum_nontrainable=0., trainable=True):
"""This is the function to perform linear interpolation between a and
b.
def lerp(a, b, momentum=0.001, momentum_nontrainable=1., trainable=True):
"""Does a linear interpolation of two parameters/ buffers.

Args:
a (float): number a
b (float): bumber b
momentum (float, optional): momentum. Defaults to 0.999.
momentum_nontrainable (float, optional): Defaults to 0.
trainable (bool, optional): trainable flag. Defaults to True.

a (torch.Tensor): Interpolation start point, refer to orig state.
b (torch.Tensor): Interpolation end point, refer to ema state.
momentum (float, optional): The weight for the interpolation
formula. Defaults to 0.001.
momentum_nontrainable (float, optional): The weight for the
interpolation formula used for nontrainable parameters.
Defaults to 1..
trainable (bool, optional): Whether input parameters is trainable.
If set to False, momentum_nontrainable will be used.
Defaults to True.
Returns:
_type_: _description_
torch.Tensor: Interpolation result.
"""
assert 0.0 < momentum < 1.0, 'momentum must be in range (0.0, 1.0)'\
f'but got {momentum}'
assert 0.0 < momentum_nontrainable <= 1.0, 'momentum_nontrainable'\
plyfager marked this conversation as resolved.
Show resolved Hide resolved
f'must be in range (0.0, 1.0] but got {momentum_nontrainable}'
m = momentum if trainable else momentum_nontrainable
return a + (b - a) * m
return b + (a - b) * m
plyfager marked this conversation as resolved.
Show resolved Hide resolved

def every_n_iters(self, runner: Runner, n: int):
"""This is the function to perform every n iterations.
Expand Down
7 changes: 5 additions & 2 deletions mmedit/models/base_models/average_model.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -230,8 +230,11 @@ def avg_func(self, averaged_param: Tensor, source_param: Tensor,
steps (int): The number of times the parameters have been
updated.
"""
momentum = self.rampup(self.steps, self.ema_kimg, self.ema_rampup,
self.batch_size, self.eps)
momentum = 1. - self.rampup(self.steps, self.ema_kimg, self.ema_rampup,
plyfager marked this conversation as resolved.
Show resolved Hide resolved
self.batch_size, self.eps)
if not (0.0 < momentum < 1.0):
warnings.warn('RampUp momentum must be in range (0.0, 1.0)'
f'but got {momentum}')
averaged_param.mul_(1 - momentum).add_(source_param, alpha=momentum)

def _load_from_state_dict(self, state_dict: dict, prefix: str,
Expand Down
2 changes: 1 addition & 1 deletion tests/test_models/test_base_models/test_average_model.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -70,7 +70,7 @@ class TestExponentialMovingAverage(TestCase):
@classmethod
def setUpClass(cls):
cls.default_cfg = dict(
interval=1, momentum=0.9999, update_buffers=True)
interval=1, momentum=0.0001, update_buffers=True)

def test_init(self):
cfg = deepcopy(self.default_cfg)
Expand Down
2 changes: 1 addition & 1 deletion tests/test_models/test_editors/test_mspie/test_mspie_stylegan2.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -29,7 +29,7 @@ def setup_class(cls):
cls.ema_config = dict(
type='ExponentialMovingAverage',
interval=1,
momentum=0.5**(32. / (ema_half_life * 1000.)))
momentum=1. - (0.5**(32. / (ema_half_life * 1000.))))

cls.loss_config = dict(
r1_loss_weight=10. / 2. * d_reg_interval,
Expand Down
2 changes: 1 addition & 1 deletion tests/test_models/test_editors/test_stylegan2/test_stylegan2.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -27,7 +27,7 @@ def setup_class(cls):
cls.ema_config = dict(
type='ExponentialMovingAverage',
interval=1,
momentum=0.5**(32. / (ema_half_life * 1000.)))
momentum=1. - (0.5**(32. / (ema_half_life * 1000.))))

cls.loss_config = dict(
r1_loss_weight=10. / 2. * d_reg_interval,
Expand Down