Python tf.test.TestCase.cached_session用法及代码示例

用法

@contextlib.contextmanager
cached_session(
    graph=None, config=None, use_gpu=True, force_gpu=False
)

返回用于执行测试的 TensorFlow 会话。

此方法的行为与 self 不同。session()：出于性能原因，cached_session 默认情况下将在同一测试中重用同一会话。此函数返回的会话只会在测试结束时关闭(在 TearDown 函数中)。

使用 use_gpu 和 force_gpu 选项来控制操作的运行位置。如果 force_gpu 为 True，则所有操作都固定到 /device:GPU:0 。否则，如果 use_gpu 为 True，TensorFlow 会尝试在 GPU 上运行尽可能多的操作。如果 force_gpu and use_gpu 都为 False，则所有操作都固定到 CPU。

例子：

class MyOperatorTest(test_util.TensorFlowTestCase):
  def testMyOperator(self):
    with self.cached_session() as sess:
      valid_input = [1.0, 2.0, 3.0, 4.0, 5.0]
      result = MyOperator(valid_input).eval()
      self.assertEqual(result, [1.0, 2.0, 3.0, 5.0, 8.0]
      invalid_input = [-1.0, 2.0, 7.0]
      with self.assertRaisesOpError("negative input not supported"):
        MyOperator(invalid_input).eval()