write_tensorboard_test.py 文件源码

def test_image_variable(self):
        """Test if ``WriteTensorBoard`` checks the image variables properly."""
        hook = WriteTensorBoard(output_dir=self.tmpdir, model=self.get_model(), image_variables=['plot'])

        with mock.patch.dict('sys.modules', **{'cv2': cv2_mock}):
            with self.assertRaises(AssertionError):
                hook.after_epoch(0, {'train': {'plot': [None]}})

            with self.assertRaises(AssertionError):
                hook.after_epoch(1, {'train': {'plot': np.zeros((10,))}})

            hook.after_epoch(2, {'train': {'plot': np.zeros((10, 10, 3))}})
        hook._summary_writer.close()