import abc import collections import contextlib import enum import itertools import re import lingvo.compat as tf from lingvo.core import cluster_factory from lingvo.core import hyperparams from lingvo.core import py_utils from tensorflow.python.ops import resource_variable_ops # pylint: disable=g-direct-tensorflow-import FLAGS = tf.flags.FLAGS _LAYER_STACK = py_utils.ThreadLocalStack() _CREATE_VARIABLES_STACK = py_utils.ThreadLocalStack() class Accumulator: """Layers can register accumulators to persist step-level state. Accumulators must be represented by a Tensor of a fixed shape. The default value must be supplied by overriding DefaultValue(). It is important that the default tensor value is created on each call in order to avoid accumulators leaking to different graphs. Accumulators can be enabled (default) or disabled by pairing Disable()/Enable() calls. When disabled, the accumulator will only return the default value and will silently drop calls to SetValue(). When computing gradients that may touch accumulators, calls should be bracketed with
# # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # ============================================================================== """Specification of a training cluster.""" import heapq import lingvo.compat as tf from lingvo.core import hyperparams from lingvo.core import py_utils import numpy as np _CLUSTER_STACK = py_utils.ThreadLocalStack() class _Cluster: """The whole training cluster from a single task's point of view.""" @classmethod def _JobSpec(cls, replicas): """Construct a job spec param with the given number of replicas.""" p = hyperparams.Params() # By default, we use /job:localhost so that most of tests can just # work out of the box. trainer.py will then set job names accordingly. p.Define('name', '/job:localhost', 'TensorFlow job spec, e.g., /job:trainer, /job:ps') p.Define('replicas', replicas, 'The number of tasks of a job.') p.Define( 'targets', '', 'The target network address(es) to which we can '
# See the License for the specific language governing permissions and # limitations under the License. # ============================================================================== """Base class for all layers.""" import abc import collections import itertools import re import lingvo.compat as tf from lingvo.core import cluster_factory from lingvo.core import hyperparams from lingvo.core import py_utils _LAYER_STACK = py_utils.ThreadLocalStack() class Accumulator: """Layers can register accumulators to persist step-level state. Accumulators must be represented by a Tensor of a fixed shape. The default value must be supplied by overriding DefaultValue(). It is important that the default tensor value is created on each call in order to avoid accumulators leaking to different graphs. Accumulators can be enabled (default) or disabled by pairing Disable()/Enable() calls. When disabled, the accumulator will only return the default value and will silently drop calls to SetValue(). When computing gradients that may touch accumulators, calls should be bracketed with Disable()/Enable().