| 227 | self.t = 1 |
| 228 | |
| 229 | def update(self, network): |
| 230 | for i, layer in enumerate(network.parametric_layers): |
| 231 | for n in layer.parameters.keys(): |
| 232 | grad = layer.parameters.grad[n] |
| 233 | self.ms[i][n] = self.beta_1 * self.ms[i][n] + (1.0 - self.beta_1) * grad |
| 234 | self.us[i][n] = np.maximum(self.beta_2 * self.us[i][n], np.abs(grad)) |
| 235 | |
| 236 | step = ( |
| 237 | self.lr |
| 238 | / (1 - self.beta_1**self.t) |
| 239 | * self.ms[i][n] |
| 240 | / (self.us[i][n] + self.epsilon) |
| 241 | ) |
| 242 | layer.parameters.step(n, -step) |
| 243 | self.t += 1 |
| 244 | |
| 245 | def setup(self, network): |
| 246 | self.ms = defaultdict(dict) |