|
|
|
from ..dist_utils import allreduce_params |
|
from .hook import HOOKS, Hook |
|
|
|
|
|
@HOOKS.register_module() |
|
class SyncBuffersHook(Hook): |
|
"""Synchronize model buffers such as running_mean and running_var in BN at |
|
the end of each epoch. |
|
|
|
Args: |
|
distributed (bool): Whether distributed training is used. It is |
|
effective only for distributed training. Defaults to True. |
|
""" |
|
|
|
def __init__(self, distributed=True): |
|
self.distributed = distributed |
|
|
|
def after_epoch(self, runner): |
|
"""All-reduce model buffers at the end of each epoch.""" |
|
if self.distributed: |
|
allreduce_params(runner.model.buffers()) |
|
|
