chore: Checking for more instances of EagerTensors
@@ -802,10 +802,20 @@ class TFKerasTrialController(det.LoopTrialController):
logging.debug(f"all-reducing logs on worker {hvd.rank()} for {len(keys)} keys {keys}.")
return {key: np.array(hvd.allreduce(logs[key], name=key)) for key in keys}
+ def _convert_possible_tensor(self, possible_tensor: Any) -> Any:
+ if isinstance(possible_tensor, EagerTensor):
+ # Horovod and / or TensorFlow may promote scalars to tensors in eager mode.
+ return possible_tensor.numpy()
+ return possible_tensor
+
def _post_train_batch_end(self, num_inputs: int, logs: Dict) -> None:
# Remove default keras metrics we aren't interested in like "batch" and "size".
self.train_workload_metrics.append(
- {k: v for k, v in logs.items() if k not in {"batch", "size"}}
+ {
+ k: self._convert_possible_tensor(v)
+ for k, v in logs.items()
+ if k not in {"batch", "size"}
+ }
)
self.train_workload_inputs += num_inputs
self.train_workload_batches += 1
@@ -820,9 +830,7 @@ class TFKerasTrialController(det.LoopTrialController):
if self.hvd_config.use:
num_inputs = hvd.allreduce(num_inputs, average=False, name="train_num_inputs")
- if isinstance(num_inputs, EagerTensor):
- # Horovod will promote an int to a tensor in eager mode.
- num_inputs = num_inputs.numpy()
+ num_inputs = self._convert_possible_tensor(num_inputs)
# Return only the latest metrics, which is the running average for all trained batches in
# the step (Keras does not report individual logs, only running averages at any point).
@@ -835,6 +843,7 @@ class TFKerasTrialController(det.LoopTrialController):
if self.is_chief:
# Don't use det.util.make_metrics, because our batch metrics are not raw metrics.
+
response = {
"metrics": {
"num_inputs": num_inputs,
