def serving_input_receiver_fn(): """Create the ServingInputReceiver to export a saved model. Returns: An instance of ServingInputReceiver. """ # We have to filter our specs since only required tensors are # used for inference time. flat_feature_spec = tensorspec_utils.flatten_spec_structure( self._get_input_features_for_receiver_fn()) required_feature_spec = ( tensorspec_utils.filter_required_flat_tensor_spec( flat_feature_spec)) receiver_tensors = tensorspec_utils.make_placeholders( required_feature_spec) # We want to ensure that our feature processing pipeline operates on a # copy of the features and does not alter the receiver_tensors. features = tensorspec_utils.flatten_spec_structure( copy.copy(receiver_tensors)) if (not self._export_raw_receivers and self._preprocess_fn is not None): features, _ = self._preprocess_fn(features=features, labels=None) return tf.estimator.export.ServingInputReceiver( features, receiver_tensors)
def map_fn(features, labels): """Creates a mapping from base_preprocessor specs to meta specs. Args: features: Features according to the spec structure of the base_preprocessor. labels: Labels according to the spec structure of the base_preprocessor. Returns: meta_features: The regrouped features according to the meta feature spec. meta_labels: The regrouped labels according to the meta label spec. """ # Now that we know all samples have the right batch size we can separate # them into data for conditioning and evaluation. features = nest.map_structure(_verify_batch_size, features) labels = nest.map_structure(_verify_batch_size, labels) condition = TSpecStructure() inference = TSpecStructure() condition.features = utils.flatten_spec_structure( nest.map_structure(_split_batch_into_condition, features)) inference.features = utils.flatten_spec_structure( nest.map_structure(_split_batch_into_inference, features)) condition.labels = utils.flatten_spec_structure( nest.map_structure(_split_batch_into_condition, labels)) meta_labels = utils.flatten_spec_structure( nest.map_structure(_split_batch_into_inference, labels)) meta_features = TSpecStructure() meta_features.condition = condition meta_features.inference = inference return meta_features, meta_labels
def create_maml_feature_spec(feature_spec, label_spec): """Create a meta feature from existing base_model specs. Note, the train spec will maintain the same name and thus mapping to the input. This is important to create the parse_tf_example_fn automatically. The validation data will have a val/ prefix such that we can feed different data to both inputs. Args: feature_spec: A hierarchy of TensorSpecs(subclasses) or Tensors. label_spec: A hierarchy of TensorSpecs(subclasses) or Tensors. Returns: An instance of TensorSpecStruct representing a valid meta learning tensor_spec with .condition and .inference access. """ condition_spec = TSpecStructure() condition_spec.features = utils.flatten_spec_structure( utils.copy_tensorspec(feature_spec, batch_size=-1, prefix='condition_features')) condition_spec.labels = utils.flatten_spec_structure( utils.copy_tensorspec(label_spec, batch_size=-1, prefix='condition_labels')) inference_spec = TSpecStructure() inference_spec.features = utils.flatten_spec_structure( utils.copy_tensorspec(feature_spec, batch_size=-1, prefix='inference_features')) meta_feature_spec = TSpecStructure() meta_feature_spec.condition = condition_spec meta_feature_spec.inference = inference_spec return meta_feature_spec
def _create_meta_spec(tensor_spec, spec_type, num_train_samples_per_task, num_val_samples_per_task): """Create a TrainValPair from an existing spec. Note, the train spec will maintain the same name and thus mapping to the input. This is important to create the parse_tf_example_fn automatically. The validation data will have a val/ prefix such that we can feed different data to both inputs. Args: tensor_spec: A dict, (named)tuple, list or a hierarchy thereof filled by TensorSpecs(subclasses) or Tensors. spec_type: A string ['features', 'labels'] specifying which spec type we alter in order to introduce the corresponding val_mode. num_train_samples_per_task: Number of training samples to expect per task batch element. num_val_samples_per_task: Number of val examples to expect per task batch element. Raises: ValueError: If the spec_type is not in ['features', 'labels']. Returns: An instance of TensorSpecStruct representing a valid meta learning tensor_spec with .train and .val access. """ if spec_type not in ['features', 'labels']: raise ValueError('We only support spec_type "features" or "labels" ' 'but received {}.'.format(spec_type)) train_tensor_spec = utils.flatten_spec_structure( utils.copy_tensorspec(tensor_spec, batch_size=num_train_samples_per_task, prefix='train')) # Since the train part is also required for inference, the specs cannot be # optional. for key, value in train_tensor_spec.items(): train_tensor_spec[key] = utils.ExtendedTensorSpec.from_spec( value, is_optional=False) val_tensor_spec = utils.flatten_spec_structure( utils.copy_tensorspec(tensor_spec, batch_size=num_val_samples_per_task, prefix='val')) # Since the train part is also required for inference, the specs for # val cannot be optional because the inputs to a while loop have to be # the same for every step of the loop. for key, value in val_tensor_spec.items(): val_tensor_spec[key] = utils.ExtendedTensorSpec.from_spec( value, is_optional=False) val_mode_shape = (1, ) if num_train_samples_per_task is None: val_mode_shape = () val_mode = TensorSpec(shape=val_mode_shape, dtype=tf.bool, name='val_mode/{}'.format(spec_type)) return utils.flatten_spec_structure( TrainValPair(train=train_tensor_spec, val=val_tensor_spec, val_mode=val_mode))
def test_pack_flat_sequence_to_spec_structure(self): subset_placeholders = utils.make_placeholders(mock_nested_subset_spec) flattened_subset_placeholders = utils.flatten_spec_structure( subset_placeholders) packed_subset_placeholders = utils.pack_flat_sequence_to_spec_structure( mock_nested_subset_spec, flattened_subset_placeholders) utils.assert_equal(subset_placeholders, packed_subset_placeholders) utils.assert_equal(mock_nested_subset_spec, packed_subset_placeholders, ignore_batch=True) placeholders = utils.make_placeholders(mock_nested_spec) flattened_placeholders = utils.flatten_spec_structure(placeholders) packed_placeholders = utils.pack_flat_sequence_to_spec_structure( mock_nested_subset_spec, flattened_placeholders) # We only subselect what we need in pack_flat_sequence_to_spec_structure, # hence, we should recover what we wanted. utils.assert_equal(mock_nested_subset_spec, packed_placeholders, ignore_batch=True) utils.assert_equal(subset_placeholders, packed_placeholders) packed_optional_placeholders = utils.pack_flat_sequence_to_spec_structure( mock_nested_optional_spec, flattened_placeholders) # Although mock_nested_optional_spec would like more tensors # flattened_placeholders cannot provide them, fortunately they are optional. utils.assert_required(packed_optional_placeholders, placeholders) utils.assert_required(mock_nested_spec, packed_optional_placeholders, ignore_batch=True)
def select_mode(val_mode, train, val): # nest.map_structure requires `native` python dicts as input since it uses # the c interface. val_dict = utils.flatten_spec_structure(val).to_dict() train_dict = utils.flatten_spec_structure(train).to_dict() select_mode_fn = lambda train, val: tf.where(val_mode, x=val, y=train) return utils.TensorSpecStruct( nest.map_structure(select_mode_fn, train_dict, val_dict).items())
def get_in_feature_specification(self, mode): feature_spec = tensorspec_utils.flatten_spec_structure( self._model_feature_specification_fn(mode)) # This will raise since we only cast from tf.float32 to tf.bfloat16. feature_spec.data_bfloat16 = tensorspec_utils.ExtendedTensorSpec.from_spec( spec=feature_spec.data_bfloat16, dtype=tf.int32) return feature_spec
def create_metaexample_spec( model_spec, num_samples_per_task, prefix): """Converts a model feature/label spec into a MetaExample spec. Args: model_spec: The base model tensor spec. num_samples_per_task: Number of episodes in the task. prefix: The tf.Example feature column name prefix. Returns: A TSpecStructure. For each spec in model_spec, the output contains num_samples_per_task corresponding specs stored as: "<name>/i". """ model_spec = utils.flatten_spec_structure(model_spec) meta_example_spec = TSpecStructure() for key in model_spec.keys(): for i in range(num_samples_per_task): spec = model_spec[key] name_prefix = '{:s}_ep{:d}'.format(prefix, i) new_name = name_prefix + '/' + six.ensure_str(spec.name) meta_example_spec[key + '/{:}'.format(i)] = ( utils.ExtendedTensorSpec.from_spec( spec, name=new_name)) return meta_example_spec
def serving_input_receiver_fn(): """Create the ServingInputReceiver to export a saved model. Returns: An instance of ServingInputReceiver. """ # We only assume one input, a string which containes the serialized proto. receiver_tensors = { 'input_example_tensor': tf.placeholder( dtype=tf.string, shape=[None], name='input_example_tensor') } feature_spec = self._get_input_features_for_receiver_fn() # We have to filter our specs since only required tensors are # used for inference time. flat_feature_spec = tensorspec_utils.flatten_spec_structure(feature_spec) required_feature_spec = ( tensorspec_utils.filter_required_flat_tensor_spec(flat_feature_spec)) tensor_dict, tensor_spec_dict = ( tensorspec_utils.tensorspec_to_feature_dict(required_feature_spec)) parse_tf_example_fn = tfdata.create_parse_tf_example_fn( tensor_dict=tensor_dict, tensor_spec_dict=tensor_spec_dict, feature_tspec=feature_spec) features = parse_tf_example_fn(receiver_tensors['input_example_tensor']) if (not self._export_raw_receivers and self._preprocess_fn is not None): features, _ = self._preprocess_fn(features=features, labels=None) return tf.estimator.export.ServingInputReceiver(features, receiver_tensors)
def test_create_metaexample_spec(self): feature_spec = TSpec() feature_spec.image = utils.ExtendedTensorSpec( shape=_DEFAULT_IN_IMAGE_SHAPE, dtype=tf.uint8, is_optional=False, data_format='jpeg', name='state/image') feature_spec.action = utils.ExtendedTensorSpec( shape=_DEFAULT_ACTION_SHAPE, dtype=tf.float32, is_optional=False, name='state/action') num_samples_in_task = 3 metaexample_spec = preprocessors.create_metaexample_spec( feature_spec, num_samples_in_task, 'condition') flat_feature_spec = utils.flatten_spec_structure(feature_spec) self.assertLen( list(metaexample_spec.keys()), num_samples_in_task * len(list(flat_feature_spec.keys()))) for key in flat_feature_spec: for i in range(num_samples_in_task): meta_example_key = six.ensure_str(key) + '/{:d}'.format(i) self.assertIn(meta_example_key, list(metaexample_spec.keys())) self.assertTrue( six.ensure_str(metaexample_spec[meta_example_key].name).startswith( 'condition_ep'))
def test_filter_required_tensor_spec_struct(self): tensor_spec_struct = utils.flatten_spec_structure( mock_nested_optional_spec) self.assertDictEqual( tensor_spec_struct, { 'train/images': T1, 'train/actions': T2, 'test/images': T1, 'test/actions': T2, 'optional/images': O4, 'optional/actions': O6, }) required_tensor_spec_struct = utils.filter_required_flat_tensor_spec( tensor_spec_struct) self.assertDictEqual( required_tensor_spec_struct, { 'train/images': T1, 'train/actions': T2, 'test/images': T1, 'test/actions': T2, }) # In case we pass a hierarchical spec this function should raise. with self.assertRaises(ValueError): required_tensor_spec_struct = utils.filter_required_flat_tensor_spec( mock_nested_optional_spec)
def test_flatten_spec_structure(self): flatten_spec_structure = utils.flatten_spec_structure( mock_nested_subset_spec) self.assertDictEqual(flatten_spec_structure, { 'train/images': T1, 'train/actions': T2 })
def predict(self, features): """Predicts based on feature input using the loaded model. Args: features: A dict containing the features used for predictions. Returns: The result of the queried model predictions. """ self.assert_is_loaded() # If using an action-tiled model, the action tiling must align with the spec # structure. If the supplied inputs align with the batch-tiled action, # expand the input to feed the tiled batch elements. flattened_feature_spec = tensorspec_utils.flatten_spec_structure( self.get_feature_specification()) def _maybe_expand_dim(path, val): model_spec = flattened_feature_spec.get(path) if model_spec and model_spec.shape.as_list() == list(val.shape): return np.expand_dims(val, 0) return val features = { k: _maybe_expand_dim(k, val) for k, val in features.items() } return self._predict_fn(features)
def create_warmup_requests_numpy(self, batch_sizes, export_dir): """Creates warm-up requests for a given feature specification. This writes an output file in `export_dir/assets.extra/tf_serving_warmup_requests` for use with Servo. Args: batch_sizes: Batch sizes of warm-up requests to write. export_dir: Base directory for the export. Returns: The filename written. """ feature_spec = self._get_input_features_for_receiver_fn() flat_feature_spec = tensorspec_utils.flatten_spec_structure( feature_spec) tf.io.gfile.makedirs(export_dir) request_filename = os.path.join(export_dir, 'tf_serving_warmup_requests') with tf.python_io.TFRecordWriter(request_filename) as writer: for batch_size in batch_sizes: request = predict_pb2.PredictRequest() request.model_spec.name = self._model_name numpy_feature_specs = tensorspec_utils.make_constant_numpy( flat_feature_spec, constant_value=0, batch_size=batch_size) for key, numpy_spec in numpy_feature_specs.items(): request.inputs[key].CopyFrom( contrib_util.make_tensor_proto(numpy_spec)) log = prediction_log_pb2.PredictionLog( predict_log=prediction_log_pb2.PredictLog(request=request)) writer.write(log.SerializeToString()) return request_filename
def test_tensor_spec_struct_composing_with_namedtuple(self): # Test that we can combine namedtuple and TensorSpecStruct. # This is important since non top level dicts use the dict_view but # nest.flatten uses the cpython interface which is why we need to maintain # a local copy of the data in subviews. test_spec = MockNested(train={ 'a': np.ones(1), 'b': 2 * np.ones(1) }, test=MockBar(images=T1, actions=T2)) ref_flat_spec = utils.flatten_spec_structure(test_spec) new_test_spec = MockNested(train=ref_flat_spec.train, test=MockBar(images=T1, actions=T2)) new_flat_spec = utils.flatten_spec_structure(new_test_spec) for key in ref_flat_spec: self.assertIn(key, new_flat_spec) self.assertEqual(new_flat_spec[key], ref_flat_spec[key])
def serving_input_receiver_fn(): """Create the ServingInputReceiver to export a saved model. Returns: An instance of ServingInputReceiver. """ # We only assume one input, a string which containes the serialized proto. receiver_tensors = { 'input_example_tensor': tf.placeholder(dtype=tf.string, shape=[None], name='input_example_tensor') } # We have to filter our specs since only required tensors are # used for inference time. flat_feature_spec = tensorspec_utils.flatten_spec_structure( self._feature_spec) # We need to freeze the conditioning and inference shapes. for key, value in flat_feature_spec.condition.items(): ref_shape = value.shape.as_list() shape = [self._num_condition_samples_per_task] + ref_shape[1:] flat_feature_spec.condition[key] = ( tensorspec_utils.ExtendedTensorSpec.from_spec(value, shape=shape)) for key, value in flat_feature_spec.inference.items(): ref_shape = value.shape.as_list() shape = [self._num_inference_samples_per_task] + ref_shape[1:] flat_feature_spec.inference[key] = ( tensorspec_utils.ExtendedTensorSpec.from_spec(value, shape=shape)) required_feature_spec = ( tensorspec_utils.filter_required_flat_tensor_spec( flat_feature_spec)) tensor_dict, tensor_spec_dict = ( tensorspec_utils.tensorspec_to_feature_dict( required_feature_spec)) parse_tf_example_fn = tfdata.create_parse_tf_example_fn( tensor_dict=tensor_dict, tensor_spec_dict=tensor_spec_dict, feature_tspec=self._feature_spec) features = parse_tf_example_fn( receiver_tensors['input_example_tensor']) if self._preprocess_fn is not None: features, _ = self._preprocess_fn( features=features, labels=None, mode=tf.estimator.ModeKeys.PREDICT) return tf.estimator.export.ServingInputReceiver( features, receiver_tensors)
def tile_val_mode(pair): """Tile val_mode to num_tasks * num_train_samples_per_task batch elements. Args: pair: TensorSpecStruct whose tensors have shape [num_tasks, num_train_samples_per_task, ...]. Raises: ValueError: If num_train_samples does not match num_val_samples. """ train_tensor = list(utils.flatten_spec_structure(pair.train).values())[0] num_train_samples_per_task = train_tensor.shape.as_list()[1] val_tensor = list(utils.flatten_spec_structure(pair.val).values())[0] num_val_samples_per_task = val_tensor.shape.as_list()[1] if num_train_samples_per_task != num_val_samples_per_task: raise ValueError('Flattening example and batch dimensions requires ' 'num_train_samples and num_val_samples to be the same.') pair.val_mode = tf.tile(pair.val_mode, [num_train_samples_per_task, 1])
def serving_input_receiver_fn(): """Create the ServingInputReceiver to export a saved model. Returns: An instance of ServingInputReceiver. """ # We have to filter our specs since only required tensors are # used for inference time. flat_feature_spec = tensorspec_utils.flatten_spec_structure( self._feature_spec) # We need to freeze the conditioning and inference shapes. for key, value in flat_feature_spec.condition.items(): ref_shape = value.shape.as_list() shape = [self._num_condition_samples_per_task] + ref_shape[1:] flat_feature_spec.condition[key] = ( tensorspec_utils.ExtendedTensorSpec.from_spec(value, shape=shape)) for key, value in flat_feature_spec.inference.items(): ref_shape = value.shape.as_list() shape = [self._num_inference_samples_per_task] + ref_shape[1:] flat_feature_spec.inference[key] = ( tensorspec_utils.ExtendedTensorSpec.from_spec(value, shape=shape)) required_feature_spec = ( tensorspec_utils.filter_required_flat_tensor_spec( flat_feature_spec)) receiver_tensors = tensorspec_utils.make_placeholders( required_feature_spec) # We want to ensure that our feature processing pipeline operates on a # copy of the features and does not alter the receiver_tensors. features = tensorspec_utils.flatten_spec_structure( copy.copy(receiver_tensors)) if self._preprocess_fn is not None: features, _ = self._preprocess_fn( features=features, labels=None, mode=tf.estimator.ModeKeys.PREDICT) return tf.estimator.export.ServingInputReceiver( features, receiver_tensors)
def get_out_label_specification(self, mode): """The specification for the output labels after executing preprocess_fn. Arguments: mode: mode key for this feature specification Returns: A TensorSpecStruct describing the required and optional tensors. """ return tensorspec_utils.flatten_spec_structure( self._model_label_specification_fn(mode))
def print_spec(tensor_spec): """Iterate over a spec and print its values in sorted order. Args: tensor_spec: A dict, (named)tuple, list or a hierarchy thereof filled by TensorSpecs(subclasses) or Tensors. """ for key, value in sorted( tensorspec_utils.flatten_spec_structure(tensor_spec).items()): logging.info('%s: %s', key, value)
def get_in_feature_specification(self, mode): """The specification for the input features for the preprocess_fn. Arguments: mode: mode key for this feature specification Returns: A TensorSpecStruct describing the required and optional tensors. """ return tensorspec_utils.flatten_spec_structure( self._model_feature_specification_fn(mode))
def get_in_feature_specification(self, mode): """See base class.""" feature_spec = tensorspec_utils.copy_tensorspec( self._model_feature_specification_fn(mode)) true_img_shape = feature_spec.image.shape.as_list() true_img_shape[ -3:-1] = self._src_img_res # Overwrite the H, W dimensions. feature_spec.image = TensorSpec.from_spec(feature_spec.image, shape=true_img_shape, dtype=tf.uint8) return tensorspec_utils.flatten_spec_structure(feature_spec)
def create_maml_label_spec(label_spec): """Create a meta feature from existing base_model specs. Args: label_spec: A hierarchy of TensorSpecs(subclasses) or Tensors. Returns: An instance of TensorSpecStruct representing a valid meta learning tensor_spec for computing the outer loss. """ return utils.flatten_spec_structure( utils.copy_tensorspec(label_spec, batch_size=-1, prefix='meta_labels'))
def infer_base_model_output_dtypes(self, mode, params): """Infer the dtypes of the model in a separate graph. Args: mode: (ModeKeys) Specifies if this is training, evaluation or prediction. params: An optional dict of hyper parameters that will be passed into input_fn and model_fn. Keys are names of parameters, values are basic python types. There are reserved keys for TPUEstimator, including 'batch_size'. Returns: dtypes: A dict containing all output dtypes {str: dtype}. """ dtype_inference_graph = tf.Graph() with dtype_inference_graph.as_default(): with tf.variable_scope('IGNORE_ONLY_TO_INFER_OUTPUT_DTYPES'): # In this graph we can now create placeholders in order to infer the # right dtype of the outputs. feature_spec = self.get_feature_specification(mode) features_dtype = utils.make_placeholders( feature_spec.condition.features, batch_size=-1) labels_dtype = utils.make_placeholders( feature_spec.condition.labels, batch_size=-1) # We need to infer the output dtypes. features_condition = utils.flatten_spec_structure( features_dtype) labels_condition = utils.flatten_spec_structure(labels_dtype) infered_outputs = self._base_model.inference_network_fn( features=features_condition, labels=labels_condition, mode=mode, params=params) dtypes = {} for key, value in infered_outputs.items(): dtypes[key] = value.dtype return dtypes
def get_in_feature_specification(self, mode): # Tempting to use create_maml_feature_spec, but we don't want the # prefixes or the meta-batch dimension. condition_spec = TSpecStructure() condition_spec.features = ( self._base_preprocessor.get_in_feature_specification(mode)) condition_spec.labels = ( self._base_preprocessor.get_in_label_specification(mode)) inference_spec = TSpecStructure() inference_spec.features = ( self._base_preprocessor.get_in_feature_specification(mode)) feature_spec = TSpecStructure() feature_spec.condition = create_metaexample_spec( condition_spec, self._num_condition_samples_per_task, 'condition') feature_spec.inference = create_metaexample_spec( inference_spec, self._num_inference_samples_per_task, 'inference') return utils.flatten_spec_structure(feature_spec)
def unflatten_batch_examples(tensor_collection, num_samples_per_task): """Unflatten task and samples dimension for M-L algorithms like RL^2. Args: tensor_collection: Tensor collection whose tensors have shape [num_tasks * num_samples_per_task, ...]. num_samples_per_task: The number of samples we have per task. Returns: result_unflattened: A tensor_collection with the same elements but the shape has been changed to [num_tasks, num_train_samples_per_task, ...]. """ result_unflattened = utils.flatten_spec_structure(tensor_collection) for key, value in result_unflattened.items(): result_unflattened[key] = tf.reshape( value, [-1, num_samples_per_task] + value.shape.as_list()[1:]) return result_unflattened
def get_in_feature_specification(self, mode ): """See base class.""" feature_spec = tensorspec_utils.copy_tensorspec( self._model_feature_specification_fn(mode)) # Don't want to parse the original_image, since we don't want to parse it # and we are adding this feature in preprocess_fn to satisfy the model's # inputs. if mode != PREDICT and 'original_image' in feature_spec: del feature_spec['original_image'] if 'image' in feature_spec: true_img_shape = feature_spec.image.shape.as_list() # Overwrite the H, W dimensions. true_img_shape[-3:-1] = self._src_img_res feature_spec.image = TensorSpec.from_spec( feature_spec.image, shape=true_img_shape, dtype=tf.uint8) return tensorspec_utils.flatten_spec_structure(feature_spec)
def _map_task_learn(self, task_learn_fn, elems, mode, params=None): """Maps a task learning/adaptation function across a batch of tasks. Args: task_learn_fn: A callable which will be mapped across `elems` once `elems` is unpacked along the first (task) dimension. elems: A (possibly) nested structure of tensors. Its tensors will be unpacked along the first dimension before passing into task_learn_fn. mode: (ModeKeys) Specifies if this is training, evaluation, or prediction. params: An optional dict of hyper parameters. Returns: The stacked outputs of task_learn_fn after being mapped across `elems`. """ if self._use_parallel_for: return pfor_map_fn(task_learn_fn, elems) else: if params is None: params = {} params['is_inner_loop'] = True dtypes = self.infer_base_model_output_dtypes(mode, params) # We flatten the features to infer the batch_size for parallel iterations. # The flattened TensorSpecStruct enables us to get the # first element of condition without knowning the name. parallel_iterations = list( utils.flatten_spec_structure( elems).values())[0].get_shape().as_list()[0] # Use parallel execution per batch, if we don't know the batch_size we # use the standard. if parallel_iterations is None: parallel_iterations = 10 # The output for val, the inner loop training steps, and training loss. dtype = ([dtypes] * 2, [dtypes] * (self._num_inner_loop_steps + 1), [tf.float32] * (self._num_inner_loop_steps + 1)) return tf.map_fn(task_learn_fn, elems=elems, dtype=dtype, parallel_iterations=parallel_iterations)
def test_is_flat_spec_or_tensors_structure(self): flatten_spec_structure = utils.flatten_spec_structure( mock_nested_subset_spec) self.assertFalse( utils.is_flat_spec_or_tensors_structure(mock_nested_subset_spec)) self.assertTrue( utils.is_flat_spec_or_tensors_structure(flatten_spec_structure)) # Flat lists are not tensor_spec_structs since they are not proper # {path: spec_or_tensor_or_numpy} structures. They would not be invariant # under flatten_spec_structure calls. self.assertFalse(utils.is_flat_spec_or_tensors_structure([T1, T2])) # Flat dictionaries are proper tensor_spec_struct_or_tensors, they would # be invariant under flatten_spec_structure. self.assertTrue( utils.is_flat_spec_or_tensors_structure({ 't1': T1, 't2': T2 }))
def get_in_feature_specification(self, mode): """The specification for the input features before executing preprocess_fn. Arguments: mode: mode key for this feature specification Returns: A TensorSpecStruct describing the required and optional tensors supporting both dictionary and hierarchical attribute like access. """ # We transform the tensorspec into it's flat structure. # This allows us to easily replace specs and replace tensors during # preprocessing. # Note, this will not alter with the spec itself, it merely allows the # user of this function to easily apply the proper preprocessing. # Further, flattening the spec structure will make a copy, allowing # us to safely operate on the in_feature_specification without # altering a implicit state or changing the model_feature_specification. tensor_spec_struct = tensorspec_utils.flatten_spec_structure( self._model_feature_specification_fn(mode)) return self._transform_in_feature_specification(tensor_spec_struct)