Module audiocraft.optim.ema

Classes

class ModuleDictEMA (module_dict: torch.nn.modules.container.ModuleDict,
decay: float = 0.999,
unbias: bool = True,
device: torch.device | str = 'cpu')

Expand source code

class ModuleDictEMA:
    """Exponential Moving Average over a nn.ModuleDict.

    You can switch to the EMA weights temporarily.
    """
    def __init__(self, module_dict: nn.ModuleDict, decay: float = 0.999,
                 unbias: bool = True, device: tp.Union[torch.device, str] = 'cpu'):
        self.decay = decay
        self.module_dict = module_dict
        self.state: dict = defaultdict(dict)
        self.count = 0
        self.device = device
        self.unbias = unbias
        self._init()

    def _init(self):
        for module_name, module in self.module_dict.items():
            for key, val in _get_named_tensors(module):
                if not val.is_floating_point():
                    continue
                device = self.device or val.device
                if key not in self.state[module_name]:
                    self.state[module_name][key] = val.detach().to(device, copy=True)

    def step(self):
        if self.unbias:
            self.count = self.count * self.decay + 1
            w = 1 / self.count
        else:
            w = 1 - self.decay
        for module_name, module in self.module_dict.items():
            for key, val in _get_named_tensors(module):
                if not val.is_floating_point():
                    continue
                device = self.device or val.device
                self.state[module_name][key].mul_(1 - w)
                self.state[module_name][key].add_(val.detach().to(device), alpha=w)

    def state_dict(self):
        return {'state': self.state, 'count': self.count}

    def load_state_dict(self, state):
        self.count = state['count']
        for module_name, module in state['state'].items():
            for key, val in module.items():
                self.state[module_name][key].copy_(val)

Exponential Moving Average over a nn.ModuleDict.

You can switch to the EMA weights temporarily.

Methods

def load_state_dict(self, state)

Expand source code

def load_state_dict(self, state):
    self.count = state['count']
    for module_name, module in state['state'].items():
        for key, val in module.items():
            self.state[module_name][key].copy_(val)

def state_dict(self)

Expand source code

def state_dict(self):
    return {'state': self.state, 'count': self.count}

def step(self)

Expand source code

def step(self):
    if self.unbias:
        self.count = self.count * self.decay + 1
        w = 1 / self.count
    else:
        w = 1 - self.decay
    for module_name, module in self.module_dict.items():
        for key, val in _get_named_tensors(module):
            if not val.is_floating_point():
                continue
            device = self.device or val.device
            self.state[module_name][key].mul_(1 - w)
            self.state[module_name][key].add_(val.detach().to(device), alpha=w)