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()