def _maybe_update_block_mask(self, weights): """Performs block-granular masking of the weights. If sparsity_m_by_n is selected, then we return the relevant pruning mask, that nullify two out of four elements in the block. Block pruning occurs only if the block_height or block_width is > 1 and if the weight tensor, when squeezed, has ndims = 2. Otherwise, elementwise pruning occurs. Args: weights: The weight tensor that needs to be masked. Returns: new_threshold: The new value of the threshold based on weights, and sparsity at the current global_step. In case of sparsity m_by_n, the returned threshold is an arbitrary number. new_mask: A numpy array of the same size and shape as weights containing 0 or 1 to indicate which of the values in weights falls below the threshold Raises: ValueError: if block pooling function is not AVG or MAX """ if self._sparsity_m_by_n: mask = self._update_mask_sparsity_m_by_n(weights, self._sparsity_m_by_n) # We need to return some numbers for threshold. return 999.0, mask if self._block_size == [1, 1]: return self._update_mask(weights) # TODO(pulkitb): Check if squeeze operations should now be removed since # we are only accepting 2-D weights. squeezed_weights = tf.squeeze(weights) abs_weights = tf.math.abs(squeezed_weights) pooled_weights = pruning_utils.factorized_pool( abs_weights, window_shape=self._block_size, pooling_type=self._block_pooling_type, strides=self._block_size, padding='SAME') if pooled_weights.get_shape().ndims != 2: pooled_weights = tf.squeeze(pooled_weights) new_threshold, new_mask = self._update_mask(pooled_weights) updated_mask = pruning_utils.expand_tensor(new_mask, self._block_size) sliced_mask = tf.slice( updated_mask, [0, 0], [squeezed_weights.get_shape()[0], squeezed_weights.get_shape()[1]]) return new_threshold, tf.reshape(sliced_mask, tf.shape(weights))
def _compare_pooling_methods(self, weights, pooling_kwargs): with self.cached_session(): variables.global_variables_initializer().run() pooled_weights_tf = array_ops.squeeze( nn_ops.pool( array_ops.reshape( weights, [1, weights.get_shape()[0], weights.get_shape()[1], 1]), **pooling_kwargs)) pooled_weights_factorized_pool = pruning_utils.factorized_pool( weights, **pooling_kwargs) self.assertAllClose(pooled_weights_tf.eval(), pooled_weights_factorized_pool.eval())
def _compare_pooling_methods(self, weights, pooling_kwargs): with self.cached_session(): compat.initialize_variables(self) pooled_weights_tf = tf.squeeze( tf.nn.pool( tf.reshape( weights, [1, weights.get_shape()[0], weights.get_shape()[1], 1]), **pooling_kwargs)) pooled_weights_factorized_pool = pruning_utils.factorized_pool( weights, **pooling_kwargs) self.assertAllClose(self.evaluate(pooled_weights_tf), self.evaluate(pooled_weights_factorized_pool))