def __init__(self,
opaque_params,
cell_nodes,
input_nodes,
fw_cell_name,
bw_cell_name,
scope=None,
name='params_canonical'):
"""Constructor.
Args:
opaque_params: opaque CuDNN params, a single tensor w/ no static shape.
cell_nodes: a int, num of nodes in a lstm cell.
input_nodes: a int, the num of nodes in input.
fw_cell_name:
bw_cell_name: the name of RNN cell in the BidiCuDNNLSTM-ish layer.
Configured via LSTMCellSimple.Params().name.
scope: the variable scope of the layer variable. If not set, default to
current variable scope.
name: name of the saveable, should be unique in a graph.
"""
self._fw_cell_name = fw_cell_name
self._bw_cell_name = bw_cell_name
scope = scope or tf.get_variable_scope()
super(CuDNNLSTMSaveable, self).__init__(opaque_params=opaque_params,
num_layers=1,
num_units=cell_nodes,
input_size=input_nodes,
direction=BI_RNN,
scope=scope,
name=name)
def CreateVariable(self, name, var_params, theta_fn=None, **kwargs):
"""Create a variable of this layer according to the parameter `var_params`.
E.g.::
def __init__(self, ...): # A layer's constructor
self.CreateVariable(
'weight', py_utils.WeightParams(shape=[100, 100]))
`theta_fn` is used to apply a simple transformation on the created
variable's value before used by the forward computation. E.g., to
add the global variational noise according to this layer's
parameter, one can do::
def __init__(self, ...): # A layer's constructor
self.CreateVariable(
name='weight',
var_params=py_utils.WeightParams(shape=[100, 100]),
theta_fn=self.AddGlobalVN)
In some contexts, eg. TPU training, variables may not be created immediately
but rather the creation request will be cached and created later via a call
to layer.CreateVariables().
Args:
name: Variable name which is used as the key into vars/theta.
var_params: `Params` used to create the variable.
theta_fn: A python function that takes a variable's value and returns a
new value to be used later for computation. Its signature must be
(tf.Tensor) -> (tf.Tensor).
**kwargs: Keyword args passed to `.py_utils.CreateVariable`.
"""
if self._is_variable_free:
raise ValueError('Cannot create variable in variable free layer.')
self._CheckName(name)
if (self.params.skip_lp_regularization and
py_utils.SKIP_LP_REGULARIZATION not in var_params.collections):
var_params = py_utils.WeightParams(
shape=var_params.shape,
dtype=var_params.dtype,
init=var_params.init,
collections=(var_params.collections +
[py_utils.SKIP_LP_REGULARIZATION]))
self._var_symbolic_shape_map[name] = var_params.shape
meta = CreateVariableMeta(
var_scope=tf.get_variable_scope(),
var_params=var_params.Copy(),
theta_fn=theta_fn,
kwargs=kwargs)
# if self._create_variables_called:
# # If CreateVariables has been called, create the variable immediately.
# self._CreateVariable(name, meta)
# else:
# # Otherwise cache the variable to be created.
# self._variables_to_create[name] = meta
self._CreateVariable(name, meta) # For now, create variables immediately.
def __init__(self, params):
super(PassiveAsymQDomain, self).__init__(params)
p = self.params
self._t_names = set() # set of known t_name (from CreateTensor)
self._qvars = py_utils.NestedMap() # var_name -> tf.Variable
# Save a scope for lazily created variables.
with tf.variable_scope(p.name + '/q'):
self._qvars_scope = tf.get_variable_scope()
def NoConstGuaranteeScope(): """Disallow const gauranteeing variable with-in scope.""" global _CONST_GUARANTEE var_scope = tf.get_variable_scope() old_caching_device = var_scope.caching_device old_val = _CONST_GUARANTEE var_scope.set_caching_device(None) _CONST_GUARANTEE = False yield _CONST_GUARANTEE = old_val var_scope.set_caching_device(old_caching_device)
def testCreateMask2D(self):
width = 10
height = 20
with self.cached_session():
weights = tf.Variable(tf.random_normal([width, height], stddev=1),
name="weights")
masked_weights = pruning.apply_mask(weights,
tf.get_variable_scope())
tf.global_variables_initializer().run()
weights_val = weights.eval()
masked_weights_val = masked_weights.eval()
self.assertAllEqual(weights_val, masked_weights_val)
def ConstGuaranteeScope(): """Treats all variables under this scope as constants.""" global _CONST_GUARANTEE var_scope = tf.get_variable_scope() old_custom_getter = var_scope.custom_getter old_caching_device = var_scope.caching_device old_val = _CONST_GUARANTEE var_scope.set_custom_getter(MaybeGuaranteeConstGetter) var_scope.set_caching_device(lambda op: op.device) _CONST_GUARANTEE = True yield _CONST_GUARANTEE = old_val var_scope.set_custom_getter(old_custom_getter) var_scope.set_caching_device(old_caching_device)
def _SelfVariableScope(self, params=None, enter_name_scope=True):
"""Internal. Used to ensure the same variable & name scopes are used."""
if not hasattr(self, '_self_variable_scope'):
params = params or self.params
self._parent_variable_scope = tf.get_variable_scope()
with tf.variable_scope(py_utils.SanitizeScopeKey(
params.name)) as scope:
self._self_variable_scope = scope
with contextlib.ExitStack() as stack:
stack.enter_context(
tf.variable_scope(self._self_variable_scope,
auxiliary_name_scope=False))
if enter_name_scope:
stack.enter_context(
tf.name_scope(
self._self_variable_scope.original_name_scope))
yield stack
def testPartitionedVariableMasking(self):
partitioner = tf.variable_axis_size_partitioner(40)
with self.cached_session() as session:
with tf.variable_scope("", partitioner=partitioner):
sparsity = tf.Variable(0.5, name="Sparsity")
weights = tf.get_variable("weights",
initializer=tf.linspace(
1.0, 100.0, 100))
masked_weights = pruning.apply_mask(
weights, scope=tf.get_variable_scope())
p = pruning.Pruning(sparsity=sparsity)
p._spec.threshold_decay = 0.0
mask_update_op = p.mask_update_op()
tf.global_variables_initializer().run()
masked_weights_val = masked_weights.eval()
session.run(mask_update_op)
masked_weights_val = masked_weights.eval()
self.assertAllEqual(np.count_nonzero(masked_weights_val), 50)
def _TestSaveRestoreHelper(self, direction):
"""Test opaque params stay 'equivalent' after save-restore."""
input_dim = 4
cell_dim = 3
with tf.variable_scope('s1'):
params_size_t = self._ParamsSize(input_dim, cell_dim, direction)
params = tf.get_variable('cudnn_params',
initializer=tf.random_uniform(
[params_size_t]),
validate_shape=False)
reset_params_op = tf.assign(params, tf.zeros_like(params))
cur_scope_name = tf.get_variable_scope().name
saveable = self._CreateSaveable(params, input_dim, cell_dim,
direction, cur_scope_name)
canonical_wts, canonical_bs = (
saveable.format_converter._opaque_to_cu_canonical(
saveable._variables))
saver = saver_lib.Saver()
with self.session(use_gpu=True) as sess:
sess.run(tf.global_variables_initializer())
save_path = os.path.join(self.get_temp_dir(), 'save-restore-unidi')
saver.save(sess, save_path)
canonical_wts_v, canonical_bs_v = sess.run(
[canonical_wts, canonical_bs])
with self.session(use_gpu=False) as sess:
sess.run(tf.global_variables_initializer())
sess.run(reset_params_op)
saver.restore(sess, save_path)
canonical_wts_v_restored, canonical_bs_v_restored = sess.run(
[canonical_wts, canonical_bs])
# Weight porition of the opaque params are exactly the same. For biases
# porition, it's expected that the sum of biases each gate stays the same.
self._CompareWeights(canonical_wts_v, canonical_wts_v_restored)
self._CompareBiases(canonical_bs_v, canonical_bs_v_restored,
direction)
def _CreateVariables(self):
# Save a scope for lazily created variables.
with tf.variable_scope('q'):
self._qvars_scope = tf.get_variable_scope()
