Sets self.optimizer. Sets self.optimizer to `optimizer`, potentially wrapping it with a LossScaleOptimizer. Args: optimizer: The optimizer(s) to assign to self.optimizer.
(self, optimizer)
| 1439 | 'found: ' + str(extra_args) + '.') |
| 1440 | |
| 1441 | def _set_optimizer(self, optimizer): |
| 1442 | """Sets self.optimizer. |
| 1443 | |
| 1444 | Sets self.optimizer to `optimizer`, potentially wrapping it with a |
| 1445 | LossScaleOptimizer. |
| 1446 | |
| 1447 | Args: |
| 1448 | optimizer: The optimizer(s) to assign to self.optimizer. |
| 1449 | """ |
| 1450 | if isinstance(optimizer, (list, tuple)): |
| 1451 | self.optimizer = [optimizers.get(opt) for opt in optimizer] |
| 1452 | else: |
| 1453 | self.optimizer = optimizers.get(optimizer) |
| 1454 | |
| 1455 | if (self._dtype_policy.loss_scale is not None and |
| 1456 | not isinstance(self.optimizer, |
| 1457 | loss_scale_optimizer.LossScaleOptimizer)): |
| 1458 | if isinstance(self.optimizer, list): |
| 1459 | raise ValueError('When a dtype policy with a loss scale is used, you ' |
| 1460 | 'can only pass a single optimizer. Using policy %s ' |
| 1461 | 'and got optimizers: %s' % |
| 1462 | self._dtype_policy, self.optimizer) |
| 1463 | if not isinstance(self.optimizer, optimizer_v2.OptimizerV2): |
| 1464 | raise ValueError('"optimizer" must be an instance of ' |
| 1465 | 'tf.keras.optimizers.Optimizer when a dype policy ' |
| 1466 | 'with a loss scale used, but got: %s. Using policy: ' |
| 1467 | '%s' % |
| 1468 | (self.optimizer, self._dtype_policy)) |
| 1469 | self.optimizer = loss_scale_optimizer.LossScaleOptimizer( |
| 1470 | self.optimizer, self._dtype_policy.loss_scale) |
| 1471 | if (isinstance(self.optimizer, loss_scale_optimizer.LossScaleOptimizer) and |
| 1472 | self._dtype_policy.loss_scale and |
| 1473 | self.optimizer.loss_scale != self._dtype_policy.loss_scale): |
| 1474 | logging.warning('LossScale of LossScaleOptimizer passed to compile (%s) ' |
| 1475 | 'is not the same as the dtype policy\'s loss scale (%s). ' |
| 1476 | 'Because the dtype policy has a loss scale, you should ' |
| 1477 | 'pass an optimizer that is not wrapped with a ' |
| 1478 | 'LossScaleOptimizer,' |
| 1479 | % (self.optimizer.loss_scale, |
| 1480 | self._dtype_policy.loss_scale)) |
| 1481 | |
| 1482 | def _prepare_validation_data(self, validation_data, batch_size, |
| 1483 | validation_steps): |