write_tensorboard_test.py 文件源码

def test_write(self):
        """Test if ``WriteTensorBoard`` writes to its FileWriter."""
        hook = WriteTensorBoard(output_dir=self.tmpdir, model=self.get_model())
        with mock.patch.object(tf.summary.FileWriter, 'add_summary') as mocked_add_summary:
            hook.after_epoch(42, {})
            self.assertEqual(mocked_add_summary.call_count, 1)
            hook.after_epoch(43, {'valid': {'accuracy': 1.0}})
            self.assertEqual(mocked_add_summary.call_count, 2)
        hook.after_epoch(44, {'valid': {'accuracy': {'mean': np.float32(1.0)}}})
        hook.after_epoch(45, {'valid': {'accuracy': {'nanmean': 1.0}}})
        hook._summary_writer.close()