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()