예제 #1
0
    def testInvalidConstructor(self):
        message = (
            'If input_tensors and output_tensors are None, both '
            'input_arrays_with_shape and output_arrays must be defined.')

        # `output_arrays` is not defined.
        with self.assertRaises(ValueError) as error:
            lite.TocoConverter(None,
                               None, [],
                               input_arrays_with_shape=[('input', [3, 9])])
        self.assertEqual(message, str(error.exception))

        # `input_arrays_with_shape` is not defined.
        with self.assertRaises(ValueError) as error:
            lite.TocoConverter(None, [], None, output_arrays=['output'])
        self.assertEqual(message, str(error.exception))
예제 #2
0
    def testValidConstructor(self):
        converter = lite.TocoConverter(None,
                                       None,
                                       None,
                                       input_arrays_with_shape=[('input',
                                                                 [3, 9])],
                                       output_arrays=['output'])
        self.assertFalse(converter._has_valid_tensors())
        self.assertEqual(converter.get_input_arrays(), ['input'])

        with self.assertRaises(ValueError) as error:
            converter._set_batch_size(1)
        self.assertEqual(
            'The batch size cannot be set for this model. Please use '
            'input_shapes parameter.', str(error.exception))

        converter = lite.TocoConverter(None, ['input_tensor'],
                                       ['output_tensor'])
        self.assertTrue(converter._has_valid_tensors())