| 1293 | |
| 1294 | |
| 1295 | class StormOptimizer(Optimizer): |
| 1296 | def __init__( |
| 1297 | self, |
| 1298 | lr=0.1, |
| 1299 | momentum=10.0, |
| 1300 | beta=0.1, |
| 1301 | grad_sq_init=0.01, |
| 1302 | policy="fixed", |
| 1303 | sparse_dedup_aggregator=None, |
| 1304 | lars=None, |
| 1305 | **kwargs |
| 1306 | ): |
| 1307 | """Constructor function to add STORM Optimizer |
| 1308 | |
| 1309 | Args: |
| 1310 | lr: learning rate scaling (called k in the original paper) |
| 1311 | momentum: momentum scaling (called c in the original paper) |
| 1312 | beta: initial value of denominator in adaptive learning rate ( |
| 1313 | called c in the original paper) |
| 1314 | grad_sq_init: initial value of gradient squared accumulator. |
| 1315 | policy: specifies how learning rate should be applied, options are |
| 1316 | 'fixed', 'step', 'exp', etc. |
| 1317 | sparse_dedup_aggregator: specifies deduplication strategy for |
| 1318 | gradient slices. Works while using sparse gradients. Options |
| 1319 | include 'mean' and 'sum'. |
| 1320 | lars: lars offset. |
| 1321 | """ |
| 1322 | super().__init__() |
| 1323 | self.lr = lr |
| 1324 | self.momentum = momentum |
| 1325 | self.beta = beta |
| 1326 | self.grad_sq_init = grad_sq_init |
| 1327 | self.policy = policy |
| 1328 | self.sparse_dedup_aggregator = sparse_dedup_aggregator |
| 1329 | self.lars = lars |
| 1330 | self.init_kwargs = kwargs |
| 1331 | |
| 1332 | def _run(self, net, param_init_net, param_info): |
| 1333 | param = param_info.blob |
| 1334 | grad = param_info.grad |
| 1335 | |
| 1336 | if self.lr <= 0: |
| 1337 | return |
| 1338 | |
| 1339 | self._clear_local_lr_multiplier() |
| 1340 | |
| 1341 | if self.lars is not None and not isinstance(grad, core.GradientSlice): |
| 1342 | assert self.lars >= 0, "Lars offset must be nonnegative, got {}".format( |
| 1343 | self.lars |
| 1344 | ) |
| 1345 | wd, trust, lr_max = self.create_lars_inputs( |
| 1346 | param_init_net, 0.0, 1.0, np.finfo(np.float32).max |
| 1347 | ) |
| 1348 | lr_lars_multiplier = net.Lars( |
| 1349 | [param, grad, wd, trust, lr_max], |
| 1350 | self.make_unique_blob_name(str(param) + "_lars"), |
| 1351 | offset=self.lars, |
| 1352 | lr_min=0.0, |
no outgoing calls
no test coverage detected
searching dependent graphs…