Home | History | Annotate | Download | only in util 
      1 # Copyright 2016 The TensorFlow Authors. All Rights Reserved.
      2 #
      3 # Licensed under the Apache License, Version 2.0 (the "License");
      4 # you may not use this file except in compliance with the License.
      5 # You may obtain a copy of the License at
      6 #
      7 #     http://www.apache.org/licenses/LICENSE-2.0
      8 #
      9 # Unless required by applicable law or agreed to in writing, software
     10 # distributed under the License is distributed on an "AS IS" BASIS,
     11 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
     12 # See the License for the specific language governing permissions and
     13 # limitations under the License.
     14 # ==============================================================================
     15 """Tests for ExampleParserConfiguration."""
     16 
     17 from __future__ import absolute_import
     18 from __future__ import division
     19 from __future__ import print_function
     20 
     21 from google.protobuf import text_format
     22 
     23 from tensorflow.core.example import example_parser_configuration_pb2
     24 from tensorflow.python.client import session
     25 from tensorflow.python.framework import dtypes
     26 from tensorflow.python.ops import array_ops
     27 from tensorflow.python.ops import parsing_ops
     28 from tensorflow.python.platform import test
     29 from tensorflow.python.util.example_parser_configuration import extract_example_parser_configuration
     30 
     31 BASIC_PROTO = """
     32 feature_map {
     33   key: "x"
     34   value {
     35     fixed_len_feature {
     36       dtype: DT_FLOAT
     37       shape {
     38         dim {
     39           size: 1
     40         }
     41       }
     42       default_value {
     43         dtype: DT_FLOAT
     44         tensor_shape {
     45           dim {
     46             size: 1
     47           }
     48         }
     49         float_val: 33.0
     50       }
     51       values_output_tensor_name: "ParseExample/ParseExample:3"
     52     }
     53   }
     54 }
     55 feature_map {
     56   key: "y"
     57   value {
     58     var_len_feature {
     59       dtype: DT_STRING
     60       values_output_tensor_name: "ParseExample/ParseExample:1"
     61       indices_output_tensor_name: "ParseExample/ParseExample:0"
     62       shapes_output_tensor_name: "ParseExample/ParseExample:2"
     63     }
     64   }
     65 }
     66 """
     67 
     68 
     69 class ExampleParserConfigurationTest(test.TestCase):
     70 
     71   def testBasic(self):
     72     golden_config = example_parser_configuration_pb2.ExampleParserConfiguration(
     73     )
     74     text_format.Parse(BASIC_PROTO, golden_config)
     75     with session.Session() as sess:
     76       examples = array_ops.placeholder(dtypes.string, shape=[1])
     77       feature_to_type = {
     78           'x': parsing_ops.FixedLenFeature([1], dtypes.float32, 33.0),
     79           'y': parsing_ops.VarLenFeature(dtypes.string)
     80       }
     81       _ = parsing_ops.parse_example(examples, feature_to_type)
     82       parse_example_op = sess.graph.get_operation_by_name(
     83           'ParseExample/ParseExample')
     84       config = extract_example_parser_configuration(parse_example_op, sess)
     85       self.assertProtoEquals(golden_config, config)
     86 
     87 
     88 if __name__ == '__main__':
     89   test.main()
     90