transform/tensorflow_transform/annotators_test.py at 3676584620d9dee74c61d5d35e46a4c95722d2b3 · smokestacklightnin/transform · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
# Copyright 2021 Google Inc. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
"""Tests for tensorflow_transform.annotators."""

import tensorflow as tf
from tensorflow_transform import annotators
from tensorflow_transform import test_case


class AnnotatorsTest(test_case.TransformTestCase):

  @test_case.named_parameters(
      dict(testcase_name='tf_compat_v1', use_tf_compat_v1=True),
      dict(testcase_name='tf2', use_tf_compat_v1=False))
  def test_annotate_asset(self, use_tf_compat_v1):
    if not use_tf_compat_v1:
      test_case.skip_if_not_tf2('Tensorflow 2.x required')

    def foo():
      annotators.annotate_asset('scope/my_key', 'scope/my_value')
      annotators.annotate_asset('my_key2', 'should_be_replaced')
      annotators.annotate_asset('my_key2', 'my_value2')

    if use_tf_compat_v1:
      with tf.Graph().as_default() as graph:
        foo()
    else:
      graph = tf.function(foo).get_concrete_function().graph

    self.assertDictEqual(
        annotators.get_asset_annotations(graph), {
            'my_key': 'my_value',
            'my_key2': 'my_value2'
        })

    annotators.clear_asset_annotations(graph)
    self.assertDictEqual(annotators.get_asset_annotations(graph), {})

  def test_object_tracker(self):
    test_case.skip_if_not_tf2('Tensorflow 2.x required')

    trackable_object = tf.__internal__.tracking.Trackable()

    @tf.function
    def preprocessing_fn():
      _ = annotators.make_and_track_object(lambda: trackable_object)
      return 1

    object_tracker = annotators.ObjectTracker()
    with annotators.object_tracker_scope(object_tracker):
      _ = preprocessing_fn()

    self.assertLen(object_tracker.trackable_objects, 1)
    self.assertEqual(trackable_object, object_tracker.trackable_objects[0])