def _update_module_A(self, module):
a = self.computeA(self.m_a[module], module)
if self.steps == 0:
self._init_A(a, module)
update_running_avg(a, self.m_A[module], self.factor_decay)
if self.sparse:
sparsification(self.m_A[module], module, ratio=self.sparse_ratio, residuals=self.residualsA)
def _update_module_G(self, module):
g = self.computeG(self.m_g[module], module, self.batch_averaged)
#logger.info('G Name: %s, shape: %s', module, g.shape)
if self.steps == 0:
self._init_G(g, module)
update_running_avg(g, self.m_G[module], self.factor_decay)
if self.sparse:
sparsification(self.m_G[module], module, ratio=self.sparse_ratio, residuals=self.residualsG)
def _update_A(self):
"""Compute and update factor A for all modules"""
for module in self.modules:
a = self.computeA(self.m_a[module], module)
if self.steps == 0:
self._init_A(a, module)
update_running_avg(a, self.m_A[module], self.factor_decay)
if self.sparse:
sparsification(self.m_A[module],
module,
ratio=self.sparse_ratio,
residuals=self.residualsA)
def _update_module_G(self, input_data, module):
g = self.computeG.get_data(input_data, module, self.batch_averaged)
if module in self.m_g:
self.m_g[module].copy_(g)
else:
self.m_g[module] = g
if self.steps == 0:
self._init_G(g, module)
#update_running_avg(g, self.m_G[module], self.factor_decay)
if self.sparse:
sparsification(self.m_G[module],
module,
ratio=self.sparse_ratio,
residuals=self.residualsG)
def _update_module_A(self, input_data, module):
a = self.computeA.get_data(input_data, module)
if module in self.m_a:
self.m_a[module].copy_(a)
else:
self.m_a[module] = a
if self.steps == 0:
self._init_A(a, module)
#update_running_avg(a, self.m_A[module], self.factor_decay)
if self.sparse:
sparsification(self.m_A[module],
module,
ratio=self.sparse_ratio,
residuals=self.residualsA)
