Ejemplo n.º 1
0
class TensorBoardReleaseManagerReleaseTest(parameterized.TestCase,
                                           unittest.IsolatedAsyncioTestCase,
                                           tf.test.TestCase):

    # pyformat: disable
    @parameterized.named_parameters(
        # materialized values
        ('bool', True, computation_types.TensorType(tf.bool), [('', True)]),
        ('int', 1, computation_types.TensorType(tf.int32), [('', 1)]),
        ('tensor_int', tf.constant(1), computation_types.TensorType(
            tf.int32), [('', tf.constant(1))]),
        ('numpy_int', np.int32(1), computation_types.TensorType(
            tf.int32), [('', np.int32(1))]),

        # value references
        ('materializable_value_reference_tensor',
         program_test_utils.TestMaterializableValueReference(1),
         computation_types.TensorType(tf.int32), [('', 1)]),

        # structures
        ('list', [
            True,
            program_test_utils.TestMaterializableValueReference(1), 'a'
        ], computation_types.SequenceType([tf.bool, tf.int32, tf.string
                                           ]), [('0', True), ('1', 1)]),
        ('list_nested', [
            [True,
             program_test_utils.TestMaterializableValueReference(1)], ['a']
        ], computation_types.SequenceType([[tf.bool, tf.int32], [tf.string]
                                           ]), [('0/0', True), ('0/1', 1)]),
        ('dict', {
            'a': True,
            'b': program_test_utils.TestMaterializableValueReference(1),
            'c': 'a'
        },
         computation_types.SequenceType([('a', tf.bool), ('b', tf.int32),
                                         ('c', tf.string)]), [('a', True),
                                                              ('b', 1)]),
        ('dict_nested', {
            'x': {
                'a': True,
                'b': program_test_utils.TestMaterializableValueReference(1)
            },
            'y': {
                'c': 'a'
            }
        },
         computation_types.SequenceType(
             [('x', [('a', tf.bool), ('b', tf.int32)]),
              ('y', [('c', tf.string)])]), [('x/a', True), ('x/b', 1)]),
        ('attr',
         program_test_utils.TestAttrObject2(
             True, program_test_utils.TestMaterializableValueReference(1)),
         computation_types.SequenceType([('a', tf.bool),
                                         ('b', tf.int32)]), [('a', True),
                                                             ('b', 1)]),
        ('attr_nested',
         program_test_utils.TestAttrObject2(
             program_test_utils.TestAttrObject2(
                 True, program_test_utils.TestMaterializableValueReference(1)),
             program_test_utils.TestAttrObject1('a')),
         computation_types.SequenceType(
             [('x', [('a', tf.bool), ('b', tf.int32)]),
              ('y', [('c', tf.string)])]), [('a/a', True), ('a/b', 1)]),
    )
    # pyformat: enable
    async def test_writes_value_scalar(self, value, type_signature,
                                       expected_calls):
        summary_dir = self.create_tempdir()
        release_mngr = tensorboard_release_manager.TensorBoardReleaseManager(
            summary_dir=summary_dir)

        with mock.patch.object(tf.summary, 'scalar') as mock_scalar:
            await release_mngr.release(value, type_signature, 1)

            self.assertLen(mock_scalar.mock_calls, len(expected_calls))
            for call, expected_args in zip(mock_scalar.mock_calls,
                                           expected_calls):
                _, args, kwargs = call
                actual_name, actual_value = args
                expected_name, expected_value = expected_args
                self.assertEqual(actual_name, expected_name)
                self.assertEqual(actual_value, expected_value)
                self.assertEqual(kwargs, {'step': 1})

    # pyformat: disable
    @parameterized.named_parameters(
        # materialized values
        ('tensor_array', tf.ones(
            [3], tf.int32), computation_types.TensorType(
                tf.int32, [3]), [('', tf.ones([3], tf.int32))]),
        ('numpy_array', np.ones(
            [3], np.int32), computation_types.TensorType(
                tf.int32, [3]), [('', np.ones([3], np.int32))]),

        # value references
        ('materializable_value_reference_sequence',
         program_test_utils.TestMaterializableValueReference(
             tf.data.Dataset.from_tensor_slices([1, 2, 3])),
         computation_types.SequenceType([tf.int32] * 3), [('', [1, 2, 3])]),
    )
    # pyformat: enable
    async def test_writes_value_histogram(self, value, type_signature,
                                          expected_calls):
        summary_dir = self.create_tempdir()
        release_mngr = tensorboard_release_manager.TensorBoardReleaseManager(
            summary_dir=summary_dir)

        with mock.patch.object(tf.summary, 'histogram') as mock_histogram:
            await release_mngr.release(value, type_signature, 1)

            self.assertLen(mock_histogram.mock_calls, len(expected_calls))
            for call, expected_args in zip(mock_histogram.mock_calls,
                                           expected_calls):
                _, args, kwargs = call
                actual_name, actual_value = args
                expected_name, expected_value = expected_args
                self.assertEqual(actual_name, expected_name)
                self.assertAllEqual(actual_value, expected_value)
                self.assertEqual(kwargs, {'step': 1})

    async def test_writes_value_scalar_and_histogram(self):
        summary_dir = self.create_tempdir()
        release_mngr = tensorboard_release_manager.TensorBoardReleaseManager(
            summary_dir=summary_dir)
        value = [1, tf.ones([3], tf.int32)]
        type_signature = computation_types.StructType([
            tf.int32,
            computation_types.TensorType(tf.float32, [3]),
        ])

        patched_scalar = mock.patch.object(tf.summary, 'scalar')
        patched_histogram = mock.patch.object(tf.summary, 'histogram')
        with patched_scalar as mock_scalar, patched_histogram as mock_histogram:
            await release_mngr.release(value, type_signature, 1)

            mock_scalar.assert_called_once_with('0', 1, step=1)
            self.assertLen(mock_histogram.mock_calls, 1)
            call = mock_histogram.mock_calls[0]
            _, args, kwargs = call
            actual_name, actual_value = args
            expected_name, expected_value = '1', tf.ones([3], tf.int32)
            self.assertEqual(actual_name, expected_name)
            self.assertAllEqual(actual_value, expected_value)
            self.assertEqual(kwargs, {'step': 1})

    @parameterized.named_parameters(
        # materialized values
        ('none', None, computation_types.StructType([])),
        ('str', 'a', computation_types.TensorType(tf.string)),
        ('tensor_str', tf.constant('a'), computation_types.TensorType(
            tf.string)),

        # structures
        ('list_empty', [], computation_types.StructType(())),
        ('dict_empty', {}, computation_types.StructType(())),
    )
    async def test_does_not_write_value(self, value, type_signature):
        summary_dir = self.create_tempdir()
        release_mngr = tensorboard_release_manager.TensorBoardReleaseManager(
            summary_dir=summary_dir)

        patch_scalar = mock.patch.object(tf.summary, 'scalar')
        patch_histogram = mock.patch.object(tf.summary, 'histogram')
        with patch_scalar as mock_scalar, patch_histogram as mock_histogram:
            await release_mngr.release(value, type_signature, 1)

            mock_scalar.assert_not_called()
            mock_histogram.assert_not_called()

    @parameterized.named_parameters(
        ('negative_1', -1),
        ('0', 0),
        ('1', 1),
    )
    async def test_does_not_raise_type_error_with_key(self, key):
        summary_dir = self.create_tempdir()
        release_mngr = tensorboard_release_manager.TensorBoardReleaseManager(
            summary_dir=summary_dir)
        value = 1
        type_signature = computation_types.TensorType(tf.int32)

        try:
            await release_mngr.release(value, type_signature, key)
        except TypeError:
            self.fail('Raised TypeError unexpectedly.')

    @parameterized.named_parameters(
        ('none', None),
        ('str', 'a'),
        ('list', []),
    )
    async def test_raises_type_error_with_key(self, key):
        summary_dir = self.create_tempdir()
        release_mngr = tensorboard_release_manager.TensorBoardReleaseManager(
            summary_dir=summary_dir)
        value = 1
        type_signature = computation_types.TensorType(tf.int32)

        with self.assertRaises(TypeError):
            await release_mngr.release(value, type_signature, key)
Ejemplo n.º 2
0
class FileProgramStateManagerLoadTest(parameterized.TestCase,
                                      unittest.IsolatedAsyncioTestCase,
                                      tf.test.TestCase):

    # pyformat: disable
    @parameterized.named_parameters(
        # materialized values
        ('none', None, None),
        ('bool', True, tf.constant(True)),
        ('int', 1, tf.constant(1)),
        ('str', 'a', tf.constant('a')),
        ('tensor_int', tf.constant(1), tf.constant(1)),
        ('tensor_str', tf.constant('a'), tf.constant('a')),
        ('tensor_array', tf.ones([3], tf.int32), tf.ones([3], tf.int32)),
        ('numpy_int', np.int32(1), tf.constant(1)),
        ('numpy_array', np.ones([3], int), tf.ones([3], tf.int32)),

        # value references
        ('materializable_value_reference_tensor',
         program_test_utils.TestMaterializableValueReference(1), tf.constant(1)
         ),
        ('materializable_value_reference_sequence',
         program_test_utils.TestMaterializableValueReference(
             tf.data.Dataset.from_tensor_slices(
                 [1, 2, 3])), tf.data.Dataset.from_tensor_slices([1, 2, 3])),

        # structures
        ('list',
         [True,
          program_test_utils.TestMaterializableValueReference(1), 'a'],
         [tf.constant(True),
          tf.constant(1), tf.constant('a')]),
        ('list_empty', [], []),
        ('list_nested', [[
            True, program_test_utils.TestMaterializableValueReference(1)
        ], ['a']], [[tf.constant(True), tf.constant(1)], [tf.constant('a')]]),
        ('dict', {
            'a': True,
            'b': program_test_utils.TestMaterializableValueReference(1),
            'c': 'a'
        }, {
            'a': tf.constant(True),
            'b': tf.constant(1),
            'c': tf.constant('a')
        }),
        ('dict_empty', {}, {}),
        ('dict_nested', {
            'x': {
                'a': True,
                'b': program_test_utils.TestMaterializableValueReference(1)
            },
            'y': {
                'c': 'a'
            }
        }, {
            'x': {
                'a': tf.constant(True),
                'b': tf.constant(1)
            },
            'y': {
                'c': tf.constant('a')
            }
        }),
        ('attr',
         program_test_utils.TestAttrObject2(
             True, program_test_utils.TestMaterializableValueReference(1)),
         program_test_utils.TestAttrObject2(tf.constant(True),
                                            tf.constant(1))),
        ('attr_nested',
         program_test_utils.TestAttrObject2(
             program_test_utils.TestAttrObject2(
                 True, program_test_utils.TestMaterializableValueReference(1)),
             program_test_utils.TestAttrObject1('a')),
         program_test_utils.TestAttrObject2(
             program_test_utils.TestAttrObject2(tf.constant(True),
                                                tf.constant(1)),
             program_test_utils.TestAttrObject1(tf.constant('a')))),
    )
    # pyformat: enable
    async def test_returns_saved_program_state(self, program_state,
                                               expected_program_state):
        root_dir = self.create_tempdir()
        program_state_mngr = file_program_state_manager.FileProgramStateManager(
            root_dir=root_dir, prefix='a_')
        await program_state_mngr.save(program_state, 1)
        structure = program_state

        actual_program_state = await program_state_mngr.load(1, structure)

        if isinstance(actual_program_state, tf.data.Dataset):
            actual_program_state = list(actual_program_state)
        if isinstance(expected_program_state, tf.data.Dataset):
            expected_program_state = list(expected_program_state)
        self.assertAllEqual(actual_program_state, expected_program_state)

    @parameterized.named_parameters(
        ('0', 0),
        ('1', 1),
        ('2', 2),
    )
    async def test_returns_saved_program_state_with_version(self, version):
        root_dir = self.create_tempdir()
        program_state_mngr = file_program_state_manager.FileProgramStateManager(
            root_dir=root_dir, prefix='a_')
        for i in range(3):
            await program_state_mngr.save(f'state_{i}', i)
        structure = 'state'

        actual_program_state = await program_state_mngr.load(
            version, structure)

        expected_program_state = f'state_{version}'
        self.assertEqual(actual_program_state, expected_program_state)

    async def test_raises_version_not_found_error_with_no_saved_program_state(
            self):
        root_dir = self.create_tempdir()
        program_state_mngr = file_program_state_manager.FileProgramStateManager(
            root_dir=root_dir, prefix='a_')

        with self.assertRaises(
                program_state_manager.ProgramStateManagerStateNotFoundError):
            await program_state_mngr.load(0, None)

    async def test_raises_version_not_found_error_with_unknown_version(self):
        root_dir = self.create_tempdir()
        program_state_mngr = file_program_state_manager.FileProgramStateManager(
            root_dir=root_dir, prefix='a_')
        await program_state_mngr.save('state_1', 1)
        structure = 'state'

        with self.assertRaises(
                program_state_manager.ProgramStateManagerStateNotFoundError):
            await program_state_mngr.load(10, structure)

    async def test_raises_structure_error(self):
        root_dir = self.create_tempdir()
        program_state_mngr = file_program_state_manager.FileProgramStateManager(
            root_dir=root_dir, prefix='a_')
        await program_state_mngr.save('state_1', 1)
        structure = []

        with self.assertRaises(
                program_state_manager.ProgramStateManagerStructureError):
            await program_state_mngr.load(1, structure)

    @parameterized.named_parameters(
        ('none', None),
        ('str', 'a'),
        ('list', []),
    )
    async def test_raises_type_error_with_version(self, version):
        root_dir = self.create_tempdir()
        program_state_mngr = file_program_state_manager.FileProgramStateManager(
            root_dir=root_dir, prefix='a_')

        with self.assertRaises(TypeError):
            await program_state_mngr.load(version, None)
Ejemplo n.º 3
0
class FileProgramStateManagerSaveTest(parameterized.TestCase,
                                      unittest.IsolatedAsyncioTestCase,
                                      tf.test.TestCase):

    # pyformat: disable
    @parameterized.named_parameters(
        # materialized values
        ('none', None, [None]),
        ('bool', True, [True]),
        ('int', 1, [1]),
        ('str', 'a', ['a']),
        ('tensor_int', tf.constant(1), [tf.constant(1)]),
        ('tensor_str', tf.constant('a'), [tf.constant('a')]),
        ('tensor_array', tf.ones([3], tf.int32), [tf.ones([3], tf.int32)]),
        ('numpy_int', np.int32(1), [np.int32(1)]),
        ('numpy_array', np.ones([3], int), [np.ones([3], int)]),

        # value references
        ('materializable_value_reference_tensor',
         program_test_utils.TestMaterializableValueReference(1), [1]),
        ('materializable_value_reference_sequence',
         program_test_utils.TestMaterializableValueReference(
             tf.data.Dataset.from_tensor_slices(
                 [1, 2, 3])), [tf.data.Dataset.from_tensor_slices([1, 2, 3])]),

        # structures
        ('list', [
            True,
            program_test_utils.TestMaterializableValueReference(1), 'a'
        ], [True, 1, 'a']),
        ('list_empty', [], []),
        ('list_nested', [[
            True, program_test_utils.TestMaterializableValueReference(1)
        ], ['a']], [True, 1, 'a']),
        ('dict', {
            'a': True,
            'b': program_test_utils.TestMaterializableValueReference(1),
            'c': 'a'
        }, [True, 1, 'a']),
        ('dict_empty', {}, []),
        ('dict_nested', {
            'x': {
                'a': True,
                'b': program_test_utils.TestMaterializableValueReference(1)
            },
            'y': {
                'c': 'a'
            }
        }, [True, 1, 'a']),
        ('attr',
         program_test_utils.TestAttrObject2(
             True, program_test_utils.TestMaterializableValueReference(1)),
         [True, 1]),
        ('attr_nested',
         program_test_utils.TestAttrObject2(
             program_test_utils.TestAttrObject2(
                 True, program_test_utils.TestMaterializableValueReference(1)),
             program_test_utils.TestAttrObject1('a')), [True, 1, 'a']),
    )
    # pyformat: enable
    async def test_writes_program_state(self, program_state, expected_value):
        root_dir = self.create_tempdir()
        program_state_mngr = file_program_state_manager.FileProgramStateManager(
            root_dir=root_dir, prefix='a_', keep_total=0)

        with mock.patch.object(file_utils,
                               'write_saved_model') as mock_write_saved_model:
            await program_state_mngr.save(program_state, 1)

            mock_write_saved_model.assert_called_once()
            call = mock_write_saved_model.mock_calls[0]
            _, args, kwargs = call
            actual_value, actual_path = args

            def _normalize(value: Any) -> Any:
                if isinstance(value, tf.data.Dataset):
                    return list(value)
                return value

            actual_value = tree.map_structure(_normalize, actual_value)
            expected_value = tree.map_structure(_normalize, expected_value)
            self.assertAllEqual(actual_value, expected_value)
            expected_path = os.path.join(root_dir, 'a_1')
            self.assertEqual(actual_path, expected_path)
            self.assertEqual(kwargs, {})

    async def test_removes_saved_program_state(self):
        root_dir = self.create_tempdir()
        program_state_mngr = file_program_state_manager.FileProgramStateManager(
            root_dir=root_dir, prefix='a_')

        with mock.patch.object(
                program_state_mngr,
                '_remove_old_program_state') as mock_remove_old_program_state:
            await program_state_mngr.save('state_1', 1)

            mock_remove_old_program_state.assert_called_once()

    async def test_raises_version_already_exists_error_with_existing_version(
            self):
        root_dir = self.create_tempdir()
        program_state_mngr = file_program_state_manager.FileProgramStateManager(
            root_dir=root_dir, prefix='a_')

        await program_state_mngr.save('state_1', 1)

        with self.assertRaises(program_state_manager.
                               ProgramStateManagerStateAlreadyExistsError):
            await program_state_mngr.save('state_1', 1)

    @parameterized.named_parameters(
        ('none', None),
        ('str', 'a'),
        ('list', []),
    )
    async def test_raises_type_error_with_version(self, version):
        root_dir = self.create_tempdir()
        program_state_mngr = file_program_state_manager.FileProgramStateManager(
            root_dir=root_dir, prefix='a_')

        with self.assertRaises(TypeError):
            await program_state_mngr.save('state', version)
Ejemplo n.º 4
0
class FlattenWithNameTest(parameterized.TestCase, tf.test.TestCase):

    # pyformat: disable
    @parameterized.named_parameters(
        # materialized values
        ('none', None, [('', None)]),
        ('bool', True, [('', True)]),
        ('int', 1, [('', 1)]),
        ('str', 'a', [('', 'a')]),
        ('tensor_int', tf.constant(1), [('', tf.constant(1))]),
        ('tensor_str', tf.constant('a'), [('', tf.constant('a'))]),
        ('tensor_array', tf.ones([3]), [('', tf.ones([3]))]),
        ('numpy_int', np.int32(1), [('', np.int32(1))]),
        ('numpy_array', np.ones([3]), [('', np.ones([3]))]),

        # value references
        ('materializable_value_reference_tensor',
         program_test_utils.TestMaterializableValueReference(1), [
             ('', program_test_utils.TestMaterializableValueReference(1))
         ]),
        ('materializable_value_reference_sequence',
         program_test_utils.TestMaterializableValueReference(
             tf.data.Dataset.from_tensor_slices([1, 2, 3])), [
                 ('',
                  program_test_utils.TestMaterializableValueReference(
                      tf.data.Dataset.from_tensor_slices([1, 2, 3])))
             ]),

        # structures
        ('list', [
            True,
            program_test_utils.TestMaterializableValueReference(1), 'a'
        ], [('0', True),
            ('1', program_test_utils.TestMaterializableValueReference(1)),
            ('2', 'a')]),
        ('list_empty', [], []),
        ('list_nested', [[
            True, program_test_utils.TestMaterializableValueReference(1)
        ], ['a']], [
            ('0/0', True),
            ('0/1', program_test_utils.TestMaterializableValueReference(1)),
            ('1/0', 'a')
        ]),
        ('dict', {
            'a': True,
            'b': program_test_utils.TestMaterializableValueReference(1),
            'c': 'a'
        }, [('a', True),
            ('b', program_test_utils.TestMaterializableValueReference(1)),
            ('c', 'a')]),
        ('dict_empty', {}, []),
        ('dict_nested', {
            'x': {
                'a': True,
                'b': program_test_utils.TestMaterializableValueReference(1)
            },
            'y': {
                'c': 'a'
            }
        }, [('x/a', True),
            ('x/b', program_test_utils.TestMaterializableValueReference(1)),
            ('y/c', 'a')]),
        ('attr',
         program_test_utils.TestAttrObject2(
             True, program_test_utils.TestMaterializableValueReference(1)), [
                 ('a', True),
                 ('b', program_test_utils.TestMaterializableValueReference(1))
             ]),
        ('attr_nested',
         program_test_utils.TestAttrObject2(
             program_test_utils.TestAttrObject2(
                 True, program_test_utils.TestMaterializableValueReference(1)),
             program_test_utils.TestAttrObject1('a')),
         [('a/a', True),
          ('a/b', program_test_utils.TestMaterializableValueReference(1)),
          ('b/a', 'a')]),
    )
    # pyformat: enable
    def test_returns_result(self, structure, expected_result):
        actual_result = structure_utils.flatten_with_name(structure)

        for actual_item, expected_item in zip(actual_result, expected_result):
            actual_path, actual_value = actual_item
            expected_path, expected_value = expected_item
            self.assertEqual(actual_path, expected_path)
            self.assertAllEqual(actual_value, expected_value)
class MemoryReleaseManagerTest(parameterized.TestCase,
                               unittest.IsolatedAsyncioTestCase,
                               tf.test.TestCase):

    # pyformat: disable
    @parameterized.named_parameters(
        # materialized values
        ('none', None, computation_types.StructType([]), None),
        ('bool', True, computation_types.TensorType(tf.bool), True),
        ('int', 1, computation_types.TensorType(tf.int32), 1),
        ('str', 'a', computation_types.TensorType(tf.string), 'a'),
        ('tensor_int', tf.constant(1), computation_types.TensorType(
            tf.int32), tf.constant(1)),
        ('tensor_str', tf.constant('a'), computation_types.TensorType(
            tf.string), tf.constant('a')),
        ('tensor_array', tf.ones([3], tf.int32),
         computation_types.TensorType(tf.int32, [3]), tf.ones([3], tf.int32)),
        ('numpy_int', np.int32(1), computation_types.TensorType(
            tf.int32), np.int32(1)),
        ('numpy_array', np.ones([3], int),
         computation_types.TensorType(tf.int32, [3]), np.ones([3], int)),

        # value references
        ('value_reference_tensor',
         program_test_utils.TestMaterializableValueReference(1),
         computation_types.TensorType(tf.int32), 1),
        ('value_reference_sequence',
         program_test_utils.TestMaterializableValueReference(
             tf.data.Dataset.from_tensor_slices(
                 [1, 2, 3])), computation_types.SequenceType(
                     tf.int32), tf.data.Dataset.from_tensor_slices([1, 2, 3])),

        # structures
        ('list', [
            True,
            program_test_utils.TestMaterializableValueReference(1), 'a'
        ], computation_types.SequenceType([tf.bool, tf.int32, tf.string
                                           ]), [True, 1, 'a']),
        ('list_empty', [], computation_types.SequenceType([]), []),
        ('list_nested', [
            [True,
             program_test_utils.TestMaterializableValueReference(1)], ['a']
        ], computation_types.SequenceType([[tf.bool, tf.int32], [tf.string]
                                           ]), [[True, 1], ['a']]),
        ('dict', {
            'a': True,
            'b': program_test_utils.TestMaterializableValueReference(1),
            'c': 'a'
        },
         computation_types.SequenceType([('a', tf.bool), ('b', tf.int32),
                                         ('c', tf.string)]), {
                                             'a': True,
                                             'b': 1,
                                             'c': 'a'
                                         }),
        ('dict_empty', {}, computation_types.SequenceType([]), {}),
        ('dict_nested', {
            'x': {
                'a': True,
                'b': program_test_utils.TestMaterializableValueReference(1)
            },
            'y': {
                'c': 'a'
            }
        },
         computation_types.SequenceType([('x', [('a', tf.bool),
                                                ('b', tf.int32)]),
                                         ('y', [('c', tf.string)])]), {
                                             'x': {
                                                 'a': True,
                                                 'b': 1
                                             },
                                             'y': {
                                                 'c': 'a'
                                             }
                                         }),
        ('attr',
         program_test_utils.TestAttrObject2(
             True, program_test_utils.TestMaterializableValueReference(1)),
         computation_types.SequenceType([
             ('a', tf.bool), ('b', tf.int32)
         ]), program_test_utils.TestAttrObject2(True, 1)),
        ('attr_nested',
         program_test_utils.TestAttrObject2(
             program_test_utils.TestAttrObject2(
                 True, program_test_utils.TestMaterializableValueReference(1)),
             program_test_utils.TestAttrObject1('a')),
         computation_types.SequenceType([('a', [('a', tf.bool),
                                                ('b', tf.int32)]),
                                         ('b', [('c', tf.string)])]),
         program_test_utils.TestAttrObject2(
             program_test_utils.TestAttrObject2(True, 1),
             program_test_utils.TestAttrObject1('a'))),
    )
    # pyformat: enable
    async def test_release_saves_value(self, value, type_signature,
                                       expected_value):
        release_mngr = memory_release_manager.MemoryReleaseManager()

        await release_mngr.release(value, type_signature, 1)

        self.assertLen(release_mngr._values, 1)
        actual_value = release_mngr._values[1]
        if isinstance(actual_value, tf.data.Dataset):
            actual_value = list(actual_value)
        if isinstance(expected_value, tf.data.Dataset):
            expected_value = list(expected_value)
        self.assertAllEqual(actual_value, expected_value)

    @parameterized.named_parameters(
        ('none', None),
        ('bool', True),
        ('int', 1),
        ('str', 'a'),
    )
    async def test_release_does_not_raise_type_error_with_key(self, key):
        release_mngr = memory_release_manager.MemoryReleaseManager()
        value = 1
        type_signature = computation_types.TensorType(tf.int32)

        try:
            await release_mngr.release(value, type_signature, key)
        except TypeError:
            self.fail('Raised TypeError unexpectedly.')

    @parameterized.named_parameters(
        ('list', []), )
    async def test_release_raises_type_error_with_key(self, key):
        release_mngr = memory_release_manager.MemoryReleaseManager()
        value = 1
        type_signature = computation_types.TensorType(tf.int32)

        with self.assertRaises(TypeError):
            await release_mngr.release(value, type_signature, key)

    @parameterized.named_parameters(
        ('0', 0),
        ('1', 1),
        ('10', 10),
    )
    def test_values_returns_values(self, count):
        release_mngr = memory_release_manager.MemoryReleaseManager()
        for i in range(count):
            release_mngr._values[i] = i * 10

        values = release_mngr.values()

        self.assertEqual(values, {i: i * 10 for i in range(count)})

    def test_values_returns_copy(self):
        release_mngr = memory_release_manager.MemoryReleaseManager()

        values_1 = release_mngr.values()
        values_2 = release_mngr.values()
        self.assertIsNot(values_1, values_2)
Ejemplo n.º 6
0
class MaterializeValueTest(parameterized.TestCase,
                           unittest.IsolatedAsyncioTestCase, tf.test.TestCase):

    # pyformat: disable
    @parameterized.named_parameters(
        # materialized values
        ('none', None, None),
        ('bool', True, True),
        ('int', 1, 1),
        ('str', 'a', 'a'),
        ('tensor_int', tf.constant(1), tf.constant(1)),
        ('tensor_str', tf.constant('a'), tf.constant('a')),
        ('tensor_array', tf.ones([3]), tf.ones([3])),
        ('numpy_int', np.int32(1), np.int32(1)),
        ('numpy_array', np.ones([3]), tf.ones([3])),

        # value references
        ('materializable_value_reference_tensor',
         program_test_utils.TestMaterializableValueReference(1), 1),
        ('materializable_value_reference_sequence',
         program_test_utils.TestMaterializableValueReference(
             tf.data.Dataset.from_tensor_slices(
                 [1, 2, 3])), tf.data.Dataset.from_tensor_slices([1, 2, 3])),

        # structures
        ('list', [
            True,
            program_test_utils.TestMaterializableValueReference(1), 'a'
        ], [True, 1, 'a']),
        ('list_empty', [], []),
        ('list_nested', [[
            True, program_test_utils.TestMaterializableValueReference(1)
        ], ['a']], [[True, 1], ['a']]),
        ('dict', {
            'a': True,
            'b': program_test_utils.TestMaterializableValueReference(1),
            'c': 'a'
        }, {
            'a': True,
            'b': 1,
            'c': 'a'
        }),
        ('dict_empty', {}, {}),
        ('dict_nested', {
            'x': {
                'a': True,
                'b': program_test_utils.TestMaterializableValueReference(1)
            },
            'y': {
                'c': 'a'
            }
        }, {
            'x': {
                'a': True,
                'b': 1
            },
            'y': {
                'c': 'a'
            }
        }),
        ('attr',
         program_test_utils.TestAttrObject2(
             True, program_test_utils.TestMaterializableValueReference(1)),
         program_test_utils.TestAttrObject2(True, 1)),
        ('attr_nested',
         program_test_utils.TestAttrObject2(
             program_test_utils.TestAttrObject2(
                 True, program_test_utils.TestMaterializableValueReference(1)),
             program_test_utils.TestAttrObject1('a')),
         program_test_utils.TestAttrObject2(
             program_test_utils.TestAttrObject2(True, 1),
             program_test_utils.TestAttrObject1('a'))),
    )
    # pyformat: enable
    async def test_returns_value(self, value, expected_value):
        actual_value = await value_reference.materialize_value(value)

        if isinstance(actual_value, tf.data.Dataset):
            actual_value = list(actual_value)
        if isinstance(expected_value, tf.data.Dataset):
            expected_value = list(expected_value)
        self.assertAllEqual(actual_value, expected_value)
class SavedModelFileReleaseManagerReleaseTest(parameterized.TestCase,
                                              unittest.IsolatedAsyncioTestCase,
                                              tf.test.TestCase):

    # pyformat: disable
    @parameterized.named_parameters(
        # materialized values
        ('none', None, computation_types.StructType([]), [None]),
        ('bool', True, computation_types.TensorType(tf.bool), [True]),
        ('int', 1, computation_types.TensorType(tf.int32), [1]),
        ('str', 'a', computation_types.TensorType(tf.string), ['a']),
        ('tensor_int', tf.constant(1), computation_types.TensorType(
            tf.int32), [tf.constant(1)]),
        ('tensor_str', tf.constant('a'), computation_types.TensorType(
            tf.string), [tf.constant('a')]),
        ('tensor_array', tf.ones([3], tf.int32),
         computation_types.TensorType(tf.int32, [3]), [tf.ones([3])]),
        ('numpy_int', np.int32(1), computation_types.TensorType(
            tf.int32), [np.int32(1)]),
        ('numpy_array', np.ones([3], int),
         computation_types.TensorType(tf.int32, [3]), [np.ones([3], int)]),

        # value references
        ('materializable_value_reference_tensor',
         program_test_utils.TestMaterializableValueReference(1),
         computation_types.TensorType(tf.int32), [1]),
        ('materializable_value_reference_sequence',
         program_test_utils.TestMaterializableValueReference(
             tf.data.Dataset.from_tensor_slices([1, 2, 3])),
         computation_types.SequenceType(
             tf.int32), [tf.data.Dataset.from_tensor_slices([1, 2, 3])]),

        # structures
        ('list', [
            True,
            program_test_utils.TestMaterializableValueReference(1), 'a'
        ], computation_types.SequenceType([tf.bool, tf.int32, tf.string
                                           ]), [True, 1, 'a']),
        ('list_empty', [], computation_types.SequenceType([]), []),
        ('list_nested', [
            [True,
             program_test_utils.TestMaterializableValueReference(1)], ['a']
        ], computation_types.SequenceType([[tf.bool, tf.int32], [tf.string]
                                           ]), [True, 1, 'a']),
        ('dict', {
            'a': True,
            'b': program_test_utils.TestMaterializableValueReference(1),
            'c': 'a'
        },
         computation_types.SequenceType([('a', tf.bool), ('b', tf.int32),
                                         ('c', tf.string)]), [True, 1, 'a']),
        ('dict_empty', {}, computation_types.SequenceType([]), []),
        ('dict_nested', {
            'x': {
                'a': True,
                'b': program_test_utils.TestMaterializableValueReference(1)
            },
            'y': {
                'c': 'a'
            }
        },
         computation_types.SequenceType(
             [('x', [('a', tf.bool), ('b', tf.int32)]),
              ('y', [('c', tf.string)])]), [True, 1, 'a']),
        ('attr',
         program_test_utils.TestAttrObject2(
             True, program_test_utils.TestMaterializableValueReference(1)),
         computation_types.SequenceType([('a', tf.bool),
                                         ('b', tf.int32)]), [True, 1]),
        ('attr_nested',
         program_test_utils.TestAttrObject2(
             program_test_utils.TestAttrObject2(
                 True, program_test_utils.TestMaterializableValueReference(1)),
             program_test_utils.TestAttrObject1('a')),
         computation_types.SequenceType(
             [('a', [('a', tf.bool), ('b', tf.int32)]),
              ('b', [('c', tf.string)])]), [True, 1, 'a']),
    )
    # pyformat: enable
    async def test_writes_value(self, value, type_signature, expected_value):
        root_dir = self.create_tempdir()
        release_mngr = file_release_manager.SavedModelFileReleaseManager(
            root_dir=root_dir, prefix='a_')

        with mock.patch.object(file_utils,
                               'write_saved_model') as mock_write_saved_model:
            await release_mngr.release(value, type_signature, 1)

            mock_write_saved_model.assert_called_once()
            call = mock_write_saved_model.mock_calls[0]
            _, args, kwargs = call
            actual_value, actual_path = args

            def _normalize(value: Any) -> Any:
                if isinstance(value, tf.data.Dataset):
                    return list(value)
                return value

            actual_value = tree.map_structure(_normalize, actual_value)
            expected_value = tree.map_structure(_normalize, expected_value)
            self.assertAllEqual(actual_value, expected_value)
            expected_path = os.path.join(root_dir, 'a_1')
            self.assertEqual(actual_path, expected_path)
            self.assertEqual(kwargs, {'overwrite': True})

    @parameterized.named_parameters(
        ('negative_1', -1),
        ('0', 0),
        ('1', 1),
    )
    async def test_does_not_raise_type_error_with_key(self, key):
        root_dir = self.create_tempdir()
        release_mngr = file_release_manager.SavedModelFileReleaseManager(
            root_dir=root_dir, prefix='a_')
        value = 1
        type_signature = computation_types.TensorType(tf.int32)

        try:
            await release_mngr.release(value, type_signature, key)
        except TypeError:
            self.fail('Raised TypeError unexpectedly.')

    @parameterized.named_parameters(
        ('none', None),
        ('str', 'a'),
        ('list', []),
    )
    async def test_raises_type_error_with_key(self, key):
        root_dir = self.create_tempdir()
        release_mngr = file_release_manager.SavedModelFileReleaseManager(
            root_dir=root_dir, prefix='a_')
        value = 1
        type_signature = computation_types.TensorType(tf.int32)

        with self.assertRaises(TypeError):
            await release_mngr.release(value, type_signature, key)
class CSVFileReleaseManagerReleaseTest(parameterized.TestCase,
                                       unittest.IsolatedAsyncioTestCase):
    async def test_calls_remove_values_greater_than_with_empty_file(self):
        file_path = self.create_tempfile()
        os.remove(file_path)
        release_mngr = file_release_manager.CSVFileReleaseManager(
            file_path=file_path)
        value = {'a': 11, 'b': 21}
        type_signature = computation_types.StructType([
            ('a', tf.int32),
            ('b', tf.int32),
        ])

        with mock.patch.object(release_mngr, '_remove_values_greater_than'
                               ) as mock_remove_values_greater_than:
            await release_mngr.release(value, type_signature, 1)

            mock_remove_values_greater_than.assert_called_once_with(0)

        self.assertEqual(release_mngr._latest_key, 1)

    async def test_calls_remove_values_greater_than_with_existing_file(self):
        file_path = self.create_tempfile()
        _write_values_to_csv(file_path=file_path,
                             fieldnames=['key', 'a', 'b'],
                             values=[{
                                 'key': 1,
                                 'a': 10,
                                 'b': 20
                             }])
        release_mngr = file_release_manager.CSVFileReleaseManager(
            file_path=file_path)
        value = {'a': 11, 'b': 21}
        type_signature = computation_types.StructType([
            ('a', tf.int32),
            ('b', tf.int32),
        ])

        with mock.patch.object(release_mngr, '_remove_values_greater_than'
                               ) as mock_remove_values_greater_than:
            await release_mngr.release(value, type_signature, 1)

            mock_remove_values_greater_than.assert_called_once_with(0)

        self.assertEqual(release_mngr._latest_key, 1)

    async def test_calls_append_value(self):
        file_path = self.create_tempfile()
        os.remove(file_path)
        release_mngr = file_release_manager.CSVFileReleaseManager(
            file_path=file_path,
            save_mode=file_release_manager.CSVSaveMode.APPEND)
        value = {'a': 11, 'b': 21}
        type_signature = computation_types.StructType([
            ('a', tf.int32),
            ('b', tf.int32),
        ])

        with mock.patch.object(release_mngr,
                               '_append_value') as mock_append_value:
            await release_mngr.release(value, type_signature, 1)

            mock_append_value.assert_called_once_with({
                'key': 1,
                'a': 11,
                'b': 21
            })

        self.assertEqual(release_mngr._latest_key, 1)

    async def test_calls_write_value(self):
        file_path = self.create_tempfile()
        os.remove(file_path)
        release_mngr = file_release_manager.CSVFileReleaseManager(
            file_path=file_path,
            save_mode=file_release_manager.CSVSaveMode.WRITE)
        value = {'a': 11, 'b': 21}
        type_signature = computation_types.StructType([
            ('a', tf.int32),
            ('b', tf.int32),
        ])

        with mock.patch.object(release_mngr,
                               '_write_value') as mock_write_value:
            await release_mngr.release(value, type_signature, 1)

            mock_write_value.assert_called_once_with({
                'key': 1,
                'a': 11,
                'b': 21
            })

        self.assertEqual(release_mngr._latest_key, 1)

    # pyformat: disable
    @parameterized.named_parameters(
        # materialized values
        ('none', None, computation_types.StructType([]), [{
            'key': '1',
            '': ''
        }]),
        ('bool', True, computation_types.TensorType(tf.bool), [{
            'key': '1',
            '': 'True'
        }]),
        ('int', 1, computation_types.TensorType(tf.int32), [{
            'key': '1',
            '': '1'
        }]),
        ('str', 'a', computation_types.TensorType(tf.string), [{
            'key': '1',
            '': 'a'
        }]),
        ('tensor_int', tf.constant(1), computation_types.TensorType(
            tf.int32), [{
                'key': '1',
                '': '1'
            }]),
        ('tensor_str', tf.constant('a'), computation_types.TensorType(
            tf.string), [{
                'key': '1',
                '': 'b\'a\''
            }]),
        ('tensor_array', tf.ones([3], tf.int32),
         computation_types.TensorType(tf.int32, [3]), [{
             'key': '1',
             '': '[1, 1, 1]'
         }]),
        ('numpy_int', np.int32(1), computation_types.TensorType(
            tf.int32), [{
                'key': '1',
                '': '1'
            }]),
        ('numpy_array', np.ones([3], int),
         computation_types.TensorType(tf.int32, [3]), [{
             'key': '1',
             '': '[1, 1, 1]'
         }]),

        # value references
        ('materializable_value_reference_tensor',
         program_test_utils.TestMaterializableValueReference(1),
         computation_types.TensorType(tf.int32), [{
             'key': '1',
             '': '1'
         }]),
        ('materializable_value_reference_sequence',
         program_test_utils.TestMaterializableValueReference(
             tf.data.Dataset.from_tensor_slices([1, 2, 3])),
         computation_types.SequenceType(tf.int32), [{
             'key': '1',
             '': '[1, 2, 3]'
         }]),

        # structures
        ('list', [
            True,
            program_test_utils.TestMaterializableValueReference(1), 'a'
        ], computation_types.SequenceType([tf.bool, tf.int32, tf.string
                                           ]), [{
                                               'key': '1',
                                               '0': 'True',
                                               '1': '1',
                                               '2': 'a'
                                           }]),
        ('list_empty', [], computation_types.SequenceType([]), [{
            'key': '1'
        }]),
        ('list_nested', [
            [True,
             program_test_utils.TestMaterializableValueReference(1)], ['a']
        ], computation_types.SequenceType([[tf.bool, tf.int32], [tf.string]
                                           ]), [{
                                               'key': '1',
                                               '0/0': 'True',
                                               '0/1': '1',
                                               '1/0': 'a'
                                           }]),
        ('dict', {
            'a': True,
            'b': program_test_utils.TestMaterializableValueReference(1),
            'c': 'a'
        },
         computation_types.SequenceType([('a', tf.bool), ('b', tf.int32),
                                         ('c', tf.string)]), [{
                                             'key': '1',
                                             'a': 'True',
                                             'b': '1',
                                             'c': 'a'
                                         }]),
        ('dict_empty', {}, computation_types.SequenceType([]), [{
            'key': '1'
        }]),
        ('dict_nested', {
            'x': {
                'a': True,
                'b': program_test_utils.TestMaterializableValueReference(1)
            },
            'y': {
                'c': 'a'
            }
        },
         computation_types.SequenceType(
             [('x', [('a', tf.bool), ('b', tf.int32)]),
              ('y', [('c', tf.string)])]), [{
                  'key': '1',
                  'x/a': 'True',
                  'x/b': '1',
                  'y/c': 'a'
              }]),
        ('attr',
         program_test_utils.TestAttrObject2(
             True, program_test_utils.TestMaterializableValueReference(1)),
         computation_types.SequenceType([('a', tf.bool),
                                         ('b', tf.int32)]), [{
                                             'key': '1',
                                             'a': 'True',
                                             'b': '1'
                                         }]),
        ('attr_nested',
         program_test_utils.TestAttrObject2(
             program_test_utils.TestAttrObject2(
                 True, program_test_utils.TestMaterializableValueReference(1)),
             program_test_utils.TestAttrObject1('a')),
         computation_types.SequenceType(
             [('a', [('a', tf.bool), ('b', tf.int32)]),
              ('b', [('c', tf.string)])]), [{
                  'key': '1',
                  'a/a': 'True',
                  'a/b': '1',
                  'b/a': 'a'
              }]),
    )
    # pyformat: enable
    async def test_writes_value(self, value, type_signature, expected_value):
        file_path = self.create_tempfile()
        os.remove(file_path)
        release_mngr = file_release_manager.CSVFileReleaseManager(
            file_path=file_path)

        await release_mngr.release(value, type_signature, 1)

        _, actual_value = _read_values_from_csv(file_path)
        self.assertEqual(actual_value, expected_value)

    @parameterized.named_parameters(
        ('negative_1', -1),
        ('0', 0),
        ('1', 1),
    )
    async def test_does_not_raise_type_error_with_key(self, key):
        file_path = self.create_tempfile()
        os.remove(file_path)
        release_mngr = file_release_manager.CSVFileReleaseManager(
            file_path=file_path)
        value = 1
        type_signature = computation_types.TensorType(tf.int32)

        try:
            await release_mngr.release(value, type_signature, key)
        except TypeError:
            self.fail('Raised TypeError unexpectedly.')

    @parameterized.named_parameters(
        ('none', None),
        ('str', 'a'),
        ('list', []),
    )
    async def test_raises_type_error_with_key(self, key):
        file_path = self.create_tempfile()
        os.remove(file_path)
        release_mngr = file_release_manager.CSVFileReleaseManager(
            file_path=file_path)
        value = 1
        type_signature = computation_types.TensorType(tf.int32)

        with self.assertRaises(TypeError):
            await release_mngr.release(value, type_signature, key)