experiments/STT-tensorflow
1
0
Fork 0
Code Issues Pull Requests Projects Releases Wiki Activity

Sync remote executors before exiting the program in TPUStrategy.

Browse Source 
PiperOrigin-RevId: 304650340
Change-Id: I7bd49dba0e63a27653dc6fcd6318295e0d4ebcf3
This commit is contained in:
Ruoxin Sang 2020-04-03 10:57:10 -07:00 committed by TensorFlower Gardener
parent 9d89f59034
commit 03b74b7a2d
1 changed files with 6 additions and 0 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

6
tensorflow/python/distribute/tpu_strategy.py
View File

@ -18,6 +18,7 @@ from __future__ import absolute_import
from __future__ import division from __future__ import division
from __future__ import print_function from __future__ import print_function
import atexit
import collections import collections
import contextlib import contextlib
import copy import copy
@ -327,6 +328,11 @@ class TPUExtended(distribute_lib.StrategyExtendedV1):
self._logical_device_stack = [0] self._logical_device_stack = [0]
if context.executing_eagerly():
# In async remote eager, we want to sync the exectors before exiting the
# program.
atexit.register(context.async_wait)
# TODO(bfontain): Remove once a proper dataset API exists for prefetching # TODO(bfontain): Remove once a proper dataset API exists for prefetching
# a dataset to multiple devices exists. # a dataset to multiple devices exists.
# If value is true, this forces prefetch of data to the host's memeory rather # If value is true, this forces prefetch of data to the host's memeory rather