array_ops.constant(6),
16,
approx=layer_collection.APPROX_DIAGONAL_NAME)
-
- self.assertEqual(9, len(lc.get_blocks()))
+ lc.register_fully_connected_multi(
+ array_ops.constant(1),
+ (array_ops.constant(2), array_ops.constant(3)),
+ (array_ops.constant(4), array_ops.constant(5)))
+ lc.register_conv2d_multi(
+ params=array_ops.ones((2, 3, 4, 5)),
+ strides=[1, 1, 1, 1],
+ padding='SAME',
+ inputs=(array_ops.ones((1, 2, 3, 4)), array_ops.ones((5, 6, 7, 8))),
+ outputs=(array_ops.ones((1, 1, 1, 5)), array_ops.ones((2, 2, 2, 10))))
+ lc.register_embedding_multi(
+ array_ops.constant((1,)),
+ (array_ops.constant(2), array_ops.constant(3)),
+ (array_ops.constant(4), array_ops.constant(5)))
+
+ self.assertEqual(12, len(lc.get_blocks()))
def testRegisterBlocksMultipleRegistrations(self):
with ops.Graph().as_default():
from tensorflow.python.framework import ops
from tensorflow.python.ops import array_ops
from tensorflow.python.ops import math_ops
+from tensorflow.python.util import nest
# For blocks corresponding to convolutional layers, or any type of block where
# the parameters can be thought of as being replicated in time or space,
PI_TYPE = pi_type
+def _make_partitionedtensors_inputs(inputs):
+ """Constructs PartitionedTensor for inputs.
+
+ The purpose of this method is to package up the towers/minibatch dimension
+ of these arrays into PartitionedTensor objects.
+
+ Args:
+ inputs: a 1-D list of Tensors. Index is tower/mini-batch.
+
+ Returns:
+ A PartitionedTensor.
+ """
+ return utils.PartitionedTensor(inputs)
+
+
+def _make_partitionedtensors_grads(grads_list):
+ """Constructs PartitionedTensor for grads_list.
+
+ The purpose of this method is to package up the towers/minibatch dimension
+ of these arrays into PartitionedTensor objects.
+
+ Args:
+ grads_list: 2-D list of Tensors. First index is for source, second
+ index for tower.
+
+ Returns:
+ Tuple of PartitionedTensors, one per source.
+ """
+ return tuple(utils.PartitionedTensor(grads) for grads in grads_list)
+
+
+def _make_partitionedtensors_multi_inputs(inputs):
+ """Constructs PartitionedTensors for inputs.
+
+ The purpose of this method is to package up the towers/minibatch dimension
+ of these arrays into PartitionedTensor objects.
+
+ This version of this function is for use with FisherBlocks that deal with
+ multiple uses or time-steps. One PartitionedTensor is created for each
+ use/time-step. The FisherBlock will be responsible for concatenating
+ (or doing whatever else it wants) with the resulting lists.
+
+ Args:
+ inputs: a 2-D list of Tensors. First index is tower/mini-batch, second is
+ use/time-step.
+
+ Returns:
+ A tuple of PartitionedTensor's, one per use/time-step.
+ """
+ num_uses = len(inputs[0])
+ assert all(len(input_) == num_uses for input_ in inputs)
+
+ return tuple(utils.PartitionedTensor(input_) for input_ in zip(*inputs))
+
+
+def _make_partitionedtensors_multi_grads(grads_list):
+ """Constructs PartitionedTensors for grads_list.
+
+ The purpose of this method is to package up the towers/minibatch dimension
+ of these arrays into PartitionedTensor objects.
+
+ This version of this function is for use with FisherBlocks that deal with
+ multiple uses or time-steps. One PartitionedTensor is created for each
+ use/time-step. The FisherBlock will be responsible for concatenating
+ (or doing whatever else it wants) with the resulting lists.
+
+ Args:
+ grads_list: 3-D list of Tensors. First index is for source, second is for
+ tower, third is for use/time-step.
+
+ Returns:
+ 2-D tuple of PartitionedTensors. First index is for source, second is for
+ use/time-step.
+ """
+ num_uses = len(grads_list[0][0])
+ assert all(len(grad) == num_uses for grads in grads_list for grad in grads)
+ return tuple(tuple(utils.PartitionedTensor(grad)
+ for grad in zip(*grads)) for grads in grads_list)
+
+
def normalize_damping(damping, num_replications):
"""Normalize damping after adjusting scale by NORMALIZE_DAMPING_POWER."""
if NORMALIZE_DAMPING_POWER:
def _outputs(self):
return self.__outputs
- def _package_minibatches(self, grads_list):
- """Constructs PartitionedTensor for inputs, grads_list.
-
- The purpose of this method is to package up the towers/minibatch dimension
- of these arrays into PartitionedTensor objects.
-
- Args:
- grads_list: 2-D list of Tensors. First index is for source, second
- index for tower.
-
- Returns:
- inputs: PartitionedTensor.
- grads_list: Tuple of PartitionedTensors, one per source.
- """
- inputs = utils.PartitionedTensor(self._inputs)
- grads_list = tuple(utils.PartitionedTensor(grads) for grads in grads_list)
-
- return inputs, grads_list
-
- def _package_minibatches_multi(self, grads_list):
- """Constructs PartitionedTensors for inputs, grads_list.
-
- The purpose of this method is to package up the towers/minibatch dimension
- of these arrays into PartitionedTensor objects.
-
- This version of this function is for use with FisherBlocks that deal with
- multiple uses or time-steps. One PartitionedTensor is created for each
- use/time-step.
-
- Args:
- grads_list: 3-D tuple of Tensors. First index is for source, second
- index is for tower, third is for use/time-step.
-
- Returns:
- inputs: A tuple of PartitionedTensor's, one per use/time-step.
- grads_list: 2-D tuple of PartitionedTensors. First index is for source,
- second is for use/time-step.
- """
- # self._inputs is a 2-D tuple. First index is tower/mini-batch, second is
- # use/time-step.
- inputs = self._inputs
- num_uses = len(inputs[0])
- assert all(len(input_) == num_uses for input_ in inputs)
- assert all(len(grad) == num_uses for grads in grads_list for grad in grads)
-
- inputs = tuple(utils.PartitionedTensor(input_) for input_ in zip(*inputs))
- grads_list = tuple(tuple(utils.PartitionedTensor(grad)
- for grad in zip(*grads)) for grads in grads_list)
-
- return inputs, grads_list
-
class FullyConnectedDiagonalFB(InputOutputMultiMinibatch, FisherBlock):
"""FisherBlock for fully-connected (dense) layers using a diagonal approx.
super(FullyConnectedDiagonalFB, self).__init__(layer_collection)
def instantiate_factors(self, grads_list, damping):
- inputs, grads_list = self._package_minibatches(grads_list)
+ inputs = _make_partitionedtensors_inputs(self._inputs)
+ grads_list = _make_partitionedtensors_grads(grads_list)
self._factor = self._layer_collection.make_or_get_factor(
fisher_factors.FullyConnectedDiagonalFactor,
super(ConvDiagonalFB, self).__init__(layer_collection)
def instantiate_factors(self, grads_list, damping):
- inputs, grads_list = self._package_minibatches(grads_list)
+ inputs = _make_partitionedtensors_inputs(self._inputs)
+ grads_list = _make_partitionedtensors_grads(grads_list)
# Infer number of locations upon which convolution is applied.
self._num_locations = num_conv_locations(inputs.shape.as_list(),
class EmbeddingKFACFB(InputOutputMultiMinibatch, KroneckerProductFB):
"""K-FAC FisherBlock for embedding layers.
- This FisherBlock is similar to EmbeddingKFACFB, except that its
+ This FisherBlock is similar to FullyConnectedKFACBasicFB, except that its
input factor is approximated by a diagonal matrix. In the case that each
example references exactly one embedding, this approximation is exact.
damping: 0-D Tensor or float. 'damping' * identity is approximately added
to this FisherBlock's Fisher approximation.
"""
- inputs, grads_list = self._package_minibatches(grads_list)
+ inputs = _make_partitionedtensors_inputs(self._inputs)
+ grads_list = _make_partitionedtensors_grads(grads_list)
- self._input_factor = self._layer_collection.make_or_get_factor( #
- fisher_factors.EmbeddingInputKroneckerFactor, #
+ self._input_factor = self._layer_collection.make_or_get_factor(
+ fisher_factors.EmbeddingInputKroneckerFactor,
(inputs, self._vocab_size))
- self._output_factor = self._layer_collection.make_or_get_factor( #
- fisher_factors.FullyConnectedKroneckerFactor, #
- (grads_list,))
+ self._output_factor = self._layer_collection.make_or_get_factor(
+ fisher_factors.FullyConnectedKroneckerFactor, (grads_list,))
self._setup_damping(damping)
+class EmbeddingKFACMultiIndepFB(InputOutputMultiMinibatch, KroneckerProductFB):
+ """K-FAC FisherBlock for embedding layers used multiple times in the graph.
+
+ Similar to EmbeddingKFACFB except that this version supports multiple uses
+ of the parameter within a single model. These uses could correspond to
+ "time-steps", but they don't have to.
+
+ Does not support bias parameters.
+ """
+
+ def __init__(self, layer_collection, vocab_size):
+ """Creates a EmbeddingKFACMultiIndepFB block.
+
+ Args:
+ layer_collection: The collection of all layers in the K-FAC approximate
+ Fisher information matrix to which this FisherBlock belongs.
+ vocab_size: int. Size of vocabulary for this embedding layer.
+ """
+ self._vocab_size = vocab_size
+
+ super(EmbeddingKFACMultiIndepFB, self).__init__(layer_collection)
+
+ def instantiate_factors(self, grads_list, damping):
+ """Instantiate Kronecker Factors for this FisherBlock.
+
+ Args:
+ grads_list: List of list of list of Tensors. grads_list[i][j][k] is the
+ gradient of the loss with respect to 'outputs' from source 'i',
+ tower/mini-batch 'j', and use/time-step 'k'. Each Tensor has shape
+ [tower_minibatch_size, output_size].
+ damping: 0-D Tensor or float. 'damping' * identity is approximately added
+ to this FisherBlock's Fisher approximation.
+ """
+ inputs = self._inputs
+ self._num_uses = num_uses = len(inputs[0])
+
+ # Check that all mini-batches/towers have the same number of uses
+ assert all(len(input_) == num_uses for input_ in inputs)
+ # Do the same for grads_list
+ assert all(len(grad) == num_uses for grad in grads for grads in grads_list)
+ # Merge uses and towers/minibatches dimensions together so we can handle
+ # it using a non-multi factor.
+ inputs = nest.flatten(inputs)
+
+ # Note that we call the multi version of make_partitionedtensors only for
+ # grads_list here.
+ inputs = _make_partitionedtensors_inputs(inputs)
+ grads_list = _make_partitionedtensors_multi_grads(grads_list)
+
+ self._input_factor = self._layer_collection.make_or_get_factor(
+ fisher_factors.EmbeddingInputKroneckerFactor,
+ (inputs, self._vocab_size))
+ self._output_factor = self._layer_collection.make_or_get_factor(
+ fisher_factors.FullyConnectedMultiKF, (grads_list,))
+ self._setup_damping(damping, normalization=num_uses)
+
+ @property
+ def _renorm_coeff(self):
+ return self._num_uses
+
+
class FullyConnectedKFACBasicFB(InputOutputMultiMinibatch, KroneckerProductFB):
"""K-FAC FisherBlock for fully-connected (dense) layers.
damping: 0-D Tensor or float. 'damping' * identity is approximately added
to this FisherBlock's Fisher approximation.
"""
- inputs, grads_list = self._package_minibatches(grads_list)
+ inputs = _make_partitionedtensors_inputs(self._inputs)
+ grads_list = _make_partitionedtensors_grads(grads_list)
- self._input_factor = self._layer_collection.make_or_get_factor( #
- fisher_factors.FullyConnectedKroneckerFactor, #
+ self._input_factor = self._layer_collection.make_or_get_factor(
+ fisher_factors.FullyConnectedKroneckerFactor,
((inputs,), self._has_bias))
- self._output_factor = self._layer_collection.make_or_get_factor( #
- fisher_factors.FullyConnectedKroneckerFactor, #
+ self._output_factor = self._layer_collection.make_or_get_factor(
+ fisher_factors.FullyConnectedKroneckerFactor,
(grads_list,))
self._setup_damping(damping)
super(ConvKFCBasicFB, self).__init__(layer_collection)
def instantiate_factors(self, grads_list, damping):
- inputs, grads_list = self._package_minibatches(grads_list)
-
# Infer number of locations upon which convolution is applied.
self._num_locations = num_conv_locations(self._inputs[0].shape.as_list(),
self._strides)
+ inputs = _make_partitionedtensors_inputs(self._inputs)
+ grads_list = _make_partitionedtensors_grads(grads_list)
+
self._input_factor = self._layer_collection.make_or_get_factor(
fisher_factors.ConvInputKroneckerFactor,
(inputs, self._filter_shape, self._padding, self._strides,
class FullyConnectedMultiIndepFB(InputOutputMultiMinibatch, KroneckerProductFB):
"""FisherBlock for fully-connected layers that share parameters.
+
+ This class implements the "independence across time" approximation from the
+ following paper:
+ https://openreview.net/pdf?id=HyMTkQZAb
"""
def __init__(self, layer_collection, has_bias=False):
def instantiate_factors(self, grads_list, damping):
self._num_uses = float(len(self._inputs[0]))
- inputs, grads_list = self._package_minibatches_multi(grads_list)
+ inputs = _make_partitionedtensors_multi_inputs(self._inputs)
+ grads_list = _make_partitionedtensors_multi_grads(grads_list)
self._input_factor = self._layer_collection.make_or_get_factor(
fisher_factors.FullyConnectedMultiKF,
return self._outputs
+class ConvKFCBasicMultiIndepFB(InputOutputMultiMinibatch, KroneckerProductFB):
+ """FisherBlock for 2D convolutional layers using the basic KFC approx.
+
+ Similar to ConvKFCBasicFB except that this version supports multiple
+ uses/time-steps via a standard independence approximation. Similar to the
+ "independence across time" used in FullyConnectedMultiIndepFB but generalized
+ in the obvious way to conv layers.
+ """
+
+ def __init__(self,
+ layer_collection,
+ params,
+ padding,
+ strides=None,
+ dilation_rate=None,
+ data_format=None,
+ extract_patches_fn=None):
+ """Creates a ConvKFCBasicMultiIndepFB block.
+
+ Args:
+ layer_collection: The collection of all layers in the K-FAC approximate
+ Fisher information matrix to which this FisherBlock belongs.
+ params: The parameters (Tensor or tuple of Tensors) of this layer. If
+ kernel alone, a Tensor of shape [..spatial_filter_shape..,
+ in_channels, out_channels]. If kernel and bias, a tuple of 2 elements
+ containing the previous and a Tensor of shape [out_channels].
+ padding: str. Padding method.
+ strides: List of ints or None. Contains [..spatial_filter_strides..] if
+ 'extract_patches_fn' is compatible with tf.nn.convolution(), else
+ [1, ..spatial_filter_strides, 1].
+ dilation_rate: List of ints or None. Rate for dilation along each spatial
+ dimension if 'extract_patches_fn' is compatible with
+ tf.nn.convolution(), else [1, ..spatial_dilation_rates.., 1].
+ data_format: str or None. Format of input data.
+ extract_patches_fn: str or None. Name of function that extracts image
+ patches. One of "extract_convolution_patches", "extract_image_patches",
+ "extract_pointwise_conv2d_patches".
+ """
+ self._padding = padding
+ self._strides = maybe_tuple(strides)
+ self._dilation_rate = maybe_tuple(dilation_rate)
+ self._data_format = data_format
+ self._extract_patches_fn = extract_patches_fn
+ self._has_bias = isinstance(params, (tuple, list))
+
+ fltr = params[0] if self._has_bias else params
+ self._filter_shape = tuple(fltr.shape.as_list())
+
+ super(ConvKFCBasicMultiIndepFB, self).__init__(layer_collection)
+
+ def instantiate_factors(self, grads_list, damping):
+ # Infer number of locations upon which convolution is applied.
+ self._num_locations = num_locations = num_conv_locations(
+ self._inputs[0][0].shape.as_list(), self._strides)
+
+ # The first index is tower/minibatch, the second is use/time-step
+ inputs = self._inputs
+ self._num_uses = num_uses = len(inputs[0])
+
+ # Check that all mini-batches/towers have the same number of uses
+ assert all(len(input_) == num_uses for input_ in inputs)
+ assert all(len(grad) == num_uses for grads in grads_list for grad in grads)
+
+ # Fold uses/time-step and towers/minibatches dimensions together
+ inputs = nest.flatten(inputs)
+ # And do the same for grads_list
+ grads_list = tuple(nest.flatten(grads) for grads in grads_list)
+
+ inputs = _make_partitionedtensors_inputs(inputs)
+ grads_list = _make_partitionedtensors_grads(grads_list)
+
+ self._input_factor = self._layer_collection.make_or_get_factor(
+ fisher_factors.ConvInputKroneckerFactor,
+ (inputs, self._filter_shape, self._padding, self._strides,
+ self._dilation_rate, self._data_format, self._extract_patches_fn,
+ self._has_bias))
+ self._output_factor = self._layer_collection.make_or_get_factor(
+ fisher_factors.ConvOutputKroneckerFactor, (grads_list,))
+
+ self._setup_damping(damping, normalization=(num_locations * num_uses))
+
+ @property
+ def _renorm_coeff(self):
+ return self._num_locations * self._num_uses
+
+
class SeriesFBApproximation(enum.IntEnum):
"""See FullyConnectedSeriesFB.__init__ for description and usage."""
option1 = 1
class FullyConnectedSeriesFB(InputOutputMultiMinibatch, FisherBlock):
"""FisherBlock for fully-connected layers that share parameters across time.
- See the following preprint for details:
+ This class implements the "Option 1" and "Option 2" approximation from the
+ following paper:
https://openreview.net/pdf?id=HyMTkQZAb
See the end of the appendix of the paper for a pseudo-code of the
def instantiate_factors(self, grads_list, damping):
self._num_timesteps = len(self._inputs[0])
- inputs, grads_list = self._package_minibatches_multi(grads_list)
+ assert len(grads_list[0][0]) == self._num_timesteps
+
+ inputs = _make_partitionedtensors_multi_inputs(self._inputs)
+ grads_list = _make_partitionedtensors_multi_grads(grads_list)
self._input_factor = self._layer_collection.make_or_get_factor(
fisher_factors.FullyConnectedMultiKF, ((inputs,), self._has_bias))
APPROX_DIAGONAL_NAME: fb.ConvDiagonalFB,
}
+_EMBEDDING_APPROX_TO_BLOCK_TYPES = {
+ APPROX_KRONECKER_NAME: fb.EmbeddingKFACFB
+}
+
APPROX_KRONECKER_INDEP_NAME = "kron_indep"
APPROX_KRONECKER_SERIES_1_NAME = "kron_series_1"
APPROX_KRONECKER_SERIES_2_NAME = "kron_series_2"
option=2)
}
+_CONV2D_MULTI_APPROX_TO_BLOCK_TYPES = {
+ APPROX_KRONECKER_INDEP_NAME: fb.ConvKFCBasicMultiIndepFB
+}
+
+_EMBEDDING_MULTI_APPROX_TO_BLOCK_TYPES = {
+ APPROX_KRONECKER_INDEP_NAME: fb.EmbeddingKFACMultiIndepFB
+}
+
# Possible value for 'reuse' keyword argument. Sets 'reuse' to
# tf.get_variable_scope().reuse.
VARIABLE_SCOPE = "VARIABLE_SCOPE"
self._default_generic_approximation = APPROX_FULL_NAME
self._default_embedding_approximation = APPROX_KRONECKER_NAME
self._default_fully_connected_approximation = APPROX_KRONECKER_NAME
- self._default_convolution_2d_approximation = APPROX_KRONECKER_NAME
+ self._default_conv2d_approximation = APPROX_KRONECKER_NAME
self._default_fully_connected_multi_approximation = (
- APPROX_KRONECKER_SERIES_2_NAME)
+ APPROX_KRONECKER_INDEP_NAME)
+ self._default_conv2d_multi_approximation = (
+ APPROX_KRONECKER_INDEP_NAME)
+ self._default_embedding_multi_approximation = APPROX_KRONECKER_INDEP_NAME
self.loss_colocation_ops = {}
self._vars_to_uses = defaultdict(lambda: 0)
@property
def default_conv2d_approximation(self):
- return self._default_convolution_2d_approximation
+ return self._default_conv2d_approximation
def set_default_conv2d_approximation(self, value):
if value not in _CONV2D_APPROX_TO_BLOCK_TYPES:
raise ValueError(
"{} is not a valid approximation for 2d convolutional layers.".format(
value))
- self._default_convolution_2d_approximation = value
+ self._default_conv2d_approximation = value
@property
def default_fully_connected_multi_approximation(self):
"multi layer.".format(value))
self._default_fully_connected_multi_approximation = value
+ @property
+ def default_conv2d_multi_approximation(self):
+ return self._default_conv2d_multi_approximation
+
+ @property
+ def default_embedding_multi_approximation(self):
+ return self._default_embedding_multi_approximation
+
def register_block(self, layer_key, fisher_block, reuse=VARIABLE_SCOPE):
"""Validates and registers the layer_key associated with the fisher_block.
else:
return None
+ def _get_block_type(self, params, approx, default, approx_to_type):
+ if approx is None:
+ approx = self._get_linked_approx(params)
+ if approx is None:
+ approx = default
+
+ if approx not in approx_to_type:
+ raise ValueError("Bad value {} for approx.".format(approx))
+
+ return approx_to_type[approx], approx
+
def register_embedding(self,
params,
inputs,
outputs,
approx=None,
reuse=VARIABLE_SCOPE):
- """Registers a fully connnected layer.
+ """Registers an embedding layer.
Args:
params: Embedding matrix of shape [vocab_size, embedding_size].
into embedding matrix.
outputs: Tensor of shape [batch_size, output_size]. Outputs
produced by layer.
- approx: str. Must be "kron".
+ approx: str or None. If not None must be "kron". The Fisher
+ approximation to use. If None the default value is used. (Default: None)
reuse: bool or str. If True, reuse an existing FisherBlock. If False,
create a new FisherBlock. If "VARIABLE_SCOPE", use
tf.get_variable_scope().reuse.
KeyError: If reuse == True but no FisherBlock found for 'params'.
ValueError: If reuse == True and FisherBlock found but of the wrong type.
"""
- if approx is None:
- approx = self._get_linked_approx(params)
- if approx is None:
- approx = self.default_embedding_approximation
-
- if approx != APPROX_KRONECKER_NAME:
- raise ValueError("Bad value {} for approx.".format(approx))
+ block_type, approx = self._get_block_type(
+ params, approx, self.default_embedding_approximation,
+ _EMBEDDING_APPROX_TO_BLOCK_TYPES)
if isinstance(params, (tuple, list)):
raise ValueError("Bias not supported.")
-
vocab_size = int(params.shape[0])
block = self.register_block(
- params, fb.EmbeddingKFACFB(self, vocab_size), reuse=reuse)
+ params, block_type(self, vocab_size), reuse=reuse)
block.register_additional_minibatch(inputs, outputs)
self._add_uses(params, 1)
inputs: Tensor of shape [batch_size, input_size]. Inputs to layer.
outputs: Tensor of shape [batch_size, output_size]. Outputs
produced by layer.
- approx: str. One of "kron" or "diagonal".
+ approx: str or None. If not None must be one of "kron" or "diagonal".
+ The Fisher approximation to use. If None the default value is used.
+ (Default: None)
reuse: bool or str. If True, reuse an existing FisherBlock. If False,
create a new FisherBlock. If "VARIABLE_SCOPE", use
tf.get_variable_scope().reuse.
KeyError: If reuse == True but no FisherBlock found for 'params'.
ValueError: If reuse == True and FisherBlock found but of the wrong type.
"""
- if approx is None:
- approx = self._get_linked_approx(params)
- if approx is None:
- approx = self.default_fully_connected_approximation
- if approx not in _FULLY_CONNECTED_APPROX_TO_BLOCK_TYPES:
- raise ValueError("Bad value {} for approx.".format(approx))
+ block_type, approx = self._get_block_type(
+ params, approx, self.default_fully_connected_approximation,
+ _FULLY_CONNECTED_APPROX_TO_BLOCK_TYPES)
- block_type = _FULLY_CONNECTED_APPROX_TO_BLOCK_TYPES[approx]
has_bias = isinstance(params, (tuple, list))
-
block = self.register_block(params, block_type(self, has_bias=has_bias),
reuse=reuse)
block.register_additional_minibatch(inputs, outputs)
Output produced by layer.
data_format: str or None. Format of data.
dilations: List of 4 ints. Dilations along each dimension.
- approx: str. One of "kron" or "diagonal".
+ approx: str or None. If not None must be one of "kron" or "diagonal".
+ The Fisher approximation to use. If None the default value is used.
+ (Default: None)
reuse: bool or str. If True, reuse an existing FisherBlock. If False,
create a new FisherBlock. If "VARIABLE_SCOPE", use
tf.get_variable_scope().reuse.
ValueError: If reuse == True and FisherBlock found but of the wrong type.
"""
- if approx is None:
- approx = self._get_linked_approx(params)
- if approx is None:
- approx = self.default_conv2d_approximation
-
- if approx not in _CONV2D_APPROX_TO_BLOCK_TYPES:
- raise ValueError("Bad value {} for approx.".format(approx))
+ block_type, approx = self._get_block_type(
+ params, approx, self.default_conv2d_approximation,
+ _CONV2D_APPROX_TO_BLOCK_TYPES)
- block_type = _CONV2D_APPROX_TO_BLOCK_TYPES[approx]
+ # It feels bad to pass in configuration that has to do with the internal
+ # implementation. And then we can't use the same constructor for both
+ # anymore and are thus forced to use this ugly if-statement.
+ # TODO(b/74793309): Clean this up?
if approx == APPROX_KRONECKER_NAME:
block = self.register_block(
params,
data_format=data_format),
reuse=reuse)
else:
- raise NotImplementedError
+ raise NotImplementedError(approx)
block.register_additional_minibatch(inputs, outputs)
dilation_rate: List of ints of length len(..input_spatial_size..).
Dilations along spatial dimension.
data_format: str or None. Format of data.
- approx: str. One of "kron" or "diagonal".
+ approx: str or None. If not None must be one of "kron" or "diagonal".
+ The Fisher approximation to use. If None the default value is used.
+ (Default: None)
reuse: bool or str. If True, reuse an existing FisherBlock. If False,
create a new FisherBlock. If "VARIABLE_SCOPE", use
tf.get_variable_scope().reuse.
KeyError: If reuse == True but no FisherBlock found for 'params'.
ValueError: If reuse == True and FisherBlock found but of the wrong type.
"""
+ # TODO(b/74793309): Have this use _get_block_type like the other
+ # registration functions?
assert approx is None or approx == APPROX_KRONECKER_NAME
block = self.register_block(
rate: None or List of ints of length 2. Dilation rates in spatial
dimensions.
data_format: str or None. Format of data.
- approx: None or str. Must be "diagonal" if non-None.
+ approx: str or None. If not None must "diagonal". The Fisher
+ approximation to use. If None the default value is used. (Default: None)
reuse: bool or str. If True, reuse an existing FisherBlock. If False,
create a new FisherBlock. If "VARIABLE_SCOPE", use
tf.get_variable_scope().reuse.
KeyError: If reuse == True but no FisherBlock found for 'params'.
ValueError: If reuse == True and FisherBlock found but of the wrong type.
"""
+ # TODO(b/74793309): Have this use _get_block_type like the other
+ # registration functions?
assert approx is None or approx == APPROX_DIAGONAL_NAME
assert data_format in [None, "NHWC"]
reuse=VARIABLE_SCOPE):
"""Register a call to tf.nn.separable_conv2d().
- Note: This requires access to intermediate outputs betwee depthwise and
+ Note: This requires access to intermediate outputs between depthwise and
pointwise convolutions.
Args:
rate: None or List of ints of length 2. Dilation rate of depthwise conv2d
kernel in spatial dimensions.
data_format: str or None. Format of data.
- approx: None or str. Must be "kron" if non-None.
+ approx: str or None. If not None must be one of "kron" or "diagonal".
+ The Fisher approximation to use. If None the default value is used.
+ (Default: None)
reuse: bool or str. If True, reuse an existing FisherBlock. If False,
create a new FisherBlock. If "VARIABLE_SCOPE", use
tf.get_variable_scope().reuse.
Args:
params: Tensor or tuple of Tensors corresponding to the parameters.
batch_size: 0-D Tensor. Size of the minibatch.
- approx: str. One of "full" or "diagonal".
+ approx: str or None. It not None, must be one of "full" or "diagonal".
+ The Fisher approximation to use. If None the default value is used.
+ (Default: None)
reuse: bool or str. If True, reuse an existing FisherBlock. If False,
create a new FisherBlock. If "VARIABLE_SCOPE", use
tf.get_variable_scope().reuse.
KeyError: If reuse == True but no FisherBlock found for 'params'.
ValueError: If reuse == True and FisherBlock found but of the wrong type.
"""
+ block_type, approx = self._get_block_type(
+ params, approx, self.default_generic_approximation,
+ _GENERIC_APPROX_TO_BLOCK_TYPES)
- if approx is None:
- approx = self._get_linked_approx(params)
- if approx is None:
- approx = self.default_generic_approximation
-
- if approx not in _GENERIC_APPROX_TO_BLOCK_TYPES:
- raise ValueError("Bad value {} for approx.".format(approx))
-
- block_type = _GENERIC_APPROX_TO_BLOCK_TYPES[approx]
block = self.register_block(params, block_type(self, params), reuse=reuse)
block.register_additional_minibatch(batch_size)
this layer. Weight matrix should have shape [input_size, output_size].
Bias should have shape [output_size].
inputs: A list of tensors, each of shape [batch_size, input_size]. Inputs
- to layer. In the case of RNNs, one Tensor per time step.
+ to layer. The list indexes each use in the graph (which might
+ correspond to a "time-step" in an RNN).
outputs: A list of tensors, the same length as 'inputs', each of shape
- [batch_size, output_size]. Outputs produced by layer. In the case of
- RNNs, one Tensor per time step.
- approx: str. One of "kron_indep", "kron_series_1", or "kron_series_2".
+ [batch_size, output_size]. Outputs produced by layer. The list indexes
+ each use in the graph (which might correspond to a "time-step" in an
+ RNN). Needs to correspond with the order used in 'inputs'.
+ approx: str or None. If not None, must be of "kron_indep", "kron_series_1"
+ or "kron_series_2". The Fisher approximation to use. If None the default
+ value is used. (Default: None)
reuse: bool or str. If True, reuse an existing FisherBlock. If False,
create a new FisherBlock. If "VARIABLE_SCOPE", use
tf.get_variable_scope().reuse.
Raises:
ValueError: For improper value to 'approx'.
"""
- if approx is None:
- approx = self._get_linked_approx(params)
- if approx is None:
- approx = self.default_fully_connected_multi_approximation
- has_bias = isinstance(params, (tuple, list))
+ block_type, approx = self._get_block_type(
+ params, approx, self.default_fully_connected_multi_approximation,
+ _FULLY_CONNECTED_MULTI_APPROX_TO_BLOCK_TYPES)
# TODO(b/70283649): something along the lines of find_canonical_output
# should be added back in here (and for the other block types, arguably).
- if approx not in _FULLY_CONNECTED_MULTI_APPROX_TO_BLOCK_TYPES:
- raise ValueError("Bad value {} for approx.".format(approx))
- block_type = _FULLY_CONNECTED_MULTI_APPROX_TO_BLOCK_TYPES[approx]
-
+ has_bias = isinstance(params, (tuple, list))
block = self.register_block(params, block_type(self, has_bias=has_bias),
reuse=reuse)
block.register_additional_minibatch(inputs, outputs)
+
+ assert len(inputs) == len(outputs)
+ self._add_uses(params, len(inputs))
+
+ def register_conv2d_multi(self,
+ params,
+ strides,
+ padding,
+ inputs,
+ outputs,
+ data_format=None,
+ dilations=None,
+ approx=None,
+ reuse=VARIABLE_SCOPE):
+ """Registers convolutional layers with shared parameters.
+
+ Args:
+ params: Tensor or 2-tuple of Tensors corresponding to weight and bias of
+ this layer. Weight matrix should have shape [kernel_height,
+ kernel_width, in_channels, out_channels]. Bias should have shape
+ [out_channels].
+ strides: 1-D Tensor of length 4. Strides for convolution kernel.
+ padding: string. see tf.nn.conv2d for valid values.
+ inputs: A list of Tensors, each of shape [batch_size, height, width,
+ in_channels]. Inputs to layer. The list indexes each use in the graph
+ (which might correspond to a "time-step" in an RNN).
+ outputs: A list of Tensors, each of shape [batch_size, height, width,
+ out_channels]. Output produced by layer. The list indexes each use
+ in the graph (which might correspond to a "time-step" in an RNN).
+ Needs to correspond with the order used in 'inputs'.
+ data_format: str or None. Format of data.
+ dilations: List of 4 ints. Dilations along each dimension.
+ approx: str or None. If not None must by "kron_indep". The Fisher
+ approximation to use. If None the default value is used.
+ (Default: None)
+ reuse: bool or str. If True, reuse an existing FisherBlock. If False,
+ create a new FisherBlock. If "VARIABLE_SCOPE", use
+ tf.get_variable_scope().reuse.
+
+ Raises:
+ ValueError: For improper value to 'approx'.
+ KeyError: If reuse == True but no FisherBlock found for 'params'.
+ ValueError: If reuse == True and FisherBlock found but of the wrong type.
+ """
+ block_type, approx = self._get_block_type(
+ params, approx, self.default_conv2d_multi_approximation,
+ _CONV2D_MULTI_APPROX_TO_BLOCK_TYPES)
+
+ block = self.register_block(
+ params,
+ block_type(
+ layer_collection=self,
+ params=params,
+ padding=padding,
+ strides=strides,
+ data_format=data_format,
+ dilation_rate=dilations,
+ extract_patches_fn="extract_image_patches"),
+ reuse=reuse)
+
+ block.register_additional_minibatch(inputs, outputs)
+
+ assert len(inputs) == len(outputs)
self._add_uses(params, len(inputs))
# TODO(b/74108452): change the loss registration functions names to refer
# to "loss functions" instead of distributions. Following naming convention
# of the loss function classes themselves.
+ def register_embedding_multi(self,
+ params,
+ inputs,
+ outputs,
+ approx=None,
+ reuse=VARIABLE_SCOPE):
+ """Registers embedding layers with shared parameters.
+
+ Args:
+ params: Embedding matrix of shape [vocab_size, embedding_size].
+ inputs: A list of Tensors, each of shape [batch_size, input_size] and
+ dtype int32. Indices into embedding matrix. The list indexes each use
+ in the graph (which might correspond to a "time-step" in an RNN).
+ outputs: A list of Tensors, each of shape [batch_size, output_size].
+ Outputs produced by layer. The list indexes each use in the graph
+ (which might correspond to a "time-step" in an RNN). Needs to
+ correspond with the order used in 'inputs'.
+ approx: str or None. If not None must by "kron_indep". The Fisher
+ approximation to use. If None the default value is used.
+ (Default: None)
+ reuse: bool or str. If True, reuse an existing FisherBlock. If False,
+ create a new FisherBlock. If "VARIABLE_SCOPE", use
+ tf.get_variable_scope().reuse.
+
+ Raises:
+ ValueError: For improper value to 'approx'.
+ KeyError: If reuse == True but no FisherBlock found for 'params'.
+ ValueError: If reuse == True and FisherBlock found but of the wrong type.
+ """
+ block_type, approx = self._get_block_type(
+ params, approx, self.default_embedding_multi_approximation,
+ _EMBEDDING_MULTI_APPROX_TO_BLOCK_TYPES)
+
+ if isinstance(params, (tuple, list)):
+ raise ValueError("Bias not supported.")
+ vocab_size = int(params.shape[0])
+
+ block = self.register_block(
+ params, block_type(self, vocab_size), reuse=reuse)
+ block.register_additional_minibatch(inputs, outputs)
+
+ self._add_uses(params, len(inputs))
+
def register_categorical_predictive_distribution(self,
logits,
seed=None,
def __hash__(self):
return hash(tuple(self.tensors))
+ def __eq__(self, other):
+ if not isinstance(other, PartitionedTensor):
+ return False
+ return self.tensors == other.tensors
+
+ def __ne__(self, other):
+ return not self == other # pylint: disable=g-comparison-negation
+
def as_tensor(self, dtype=None, name=None, as_ref=False):
with ops.name_scope(name, "PartitionedTensor.as_tensor", self.tensors):
assert not as_ref
projects / platform / upstream / tensorflow.git / commitdiff