write_tensorboard_test.py 文件源码

def test_args(self):
        """Test WriteTensorBoard argument handling and ``SummaryWriter`` creation."""
        model = self.get_model()

        with self.assertRaises(AssertionError):
            _ = WriteTensorBoard(output_dir=self.tmpdir, model=42)
        with self.assertRaises(AssertionError):
            _ = WriteTensorBoard(output_dir=self.tmpdir, model=model, on_missing_variable='not-recognized')
        with self.assertRaises(AssertionError):
            _ = WriteTensorBoard(output_dir=self.tmpdir, model=model, on_unknown_type='not-recognized')

        with mock.patch('tensorflow.summary.FileWriter', autospec=True) as mocked_writer:
            _ = WriteTensorBoard(output_dir=self.tmpdir, model=model, flush_secs=42, visualize_graph=True)
            mocked_writer.assert_called_with(logdir=self.tmpdir, flush_secs=42, graph=model.graph)

            _ = WriteTensorBoard(output_dir=self.tmpdir, model=model)
            mocked_writer.assert_called_with(logdir=self.tmpdir, flush_secs=10, graph=None)