Python context.device函数代码示例

  def disabled_testSmallIntegerOpsForcedToCPU(self):
    if not context.context().num_gpus():
      self.skipTest('No GPUs found')

    a = constant_op.constant((1, 2, 3, 4, 5), dtype=dtypes.int64)
    b = constant_op.constant((2, 3, 4, 5, 6), dtype=dtypes.int64)
    with context.device('gpu:0'):
      c = a + b

    # Op forced to CPU since all constants are integers and small.
    self.assertEqual(c.device, '/job:localhost/replica:0/task:0/device:CPU:0')

    a = array_ops.zeros((8, 10), dtype=dtypes.int64)
    b = array_ops.ones((8, 10), dtype=dtypes.int64)

    with context.device('gpu:0'):
      c = a + b

    # Op not forced to CPU since the tensors are larger than 64 elements.
    self.assertEqual(c.device, '/job:localhost/replica:0/task:0/device:GPU:0')

    a = constant_op.constant((1, 2, 3, 4, 5), dtype=dtypes.float32)
    b = constant_op.constant((2, 3, 4, 5, 6), dtype=dtypes.float32)
    with context.device('gpu:0'):
      c = a + b

    # Op not forced to CPU since the constants are not integers.
    self.assertEqual(c.device, '/job:localhost/replica:0/task:0/device:GPU:0')

 def testResourceTensorPlacement(self):
   with context.device('gpu:0'):
     v = resource_variable_ops.ResourceVariable(1.0)
   with context.device('cpu:0'):
     # Check that even though we specified the cpu device we'll run the read op
     # in the device where the handle is.
     self.assertAllEqual(
         gen_resource_variable_ops.read_variable_op(v.handle, v.dtype), 1.0)

 def testGenericSummary(self):
   with context.device(self._test_device):
     x = constant_op.constant(1337.0)
     with context.device("cpu:0"):
       metadata = constant_op.constant("foo")
     self._writer.generic("x", x, metadata)
     event = self._readLastEvent()
     self.assertEqual("x", event.summary.value[0].tag)

 def testReEntrant(self):
   cpu = context.device('cpu:0')
   gpu = context.device('gpu:0')
   with cpu:
     with gpu:
       with gpu:
         self.assertEndsWith(current_device(), 'GPU:0')
       self.assertEndsWith(current_device(), 'GPU:0')
     self.assertEndsWith(current_device(), 'CPU:0')
     with gpu:
       self.assertEndsWith(current_device(), 'GPU:0')

 def testDevicePlacementEnforcesConsistency(self):
   cpu = context.device('cpu:0')
   gpu = context.device('gpu:0')
   cpu.__enter__()
   self.assertEndsWith(current_device(), 'CPU:0')
   gpu.__enter__()
   self.assertEndsWith(current_device(), 'GPU:0')
   with self.assertRaisesRegexp(
       RuntimeError, 'Exiting device scope without proper scope nesting'):
     cpu.__exit__()
     self.assertEndsWith(current_device(), 'GPU:0')
   gpu.__exit__()
   self.assertEndsWith(current_device(), 'CPU:0')

 def run_eager_mode(self, **kwargs):
   if force_gpu:
     gpu_name = gpu_device_name()
     if not gpu_name:
       gpu_name = "/device:GPU:0"
     with context.device(gpu_name):
       f(self)
   elif use_gpu:
     # TODO(xpan): Support softplacement and gpu by default when available.
     f(self, **kwargs)
   else:
     with context.device("/device:CPU:0"):
       f(self, **kwargs)

  def testReEntrant(self):
    if not context.context().num_gpus():
      self.skipTest('No GPUs found')

    cpu = context.device('cpu:0')
    gpu = context.device('gpu:0')
    with cpu:
      with gpu:
        with gpu:
          self.assertEndsWith(current_device(), 'GPU:0')
        self.assertEndsWith(current_device(), 'GPU:0')
      self.assertEndsWith(current_device(), 'CPU:0')
      with gpu:
        self.assertEndsWith(current_device(), 'GPU:0')

 def testImageSummary(self):
   with context.device(self._test_device):
     a = constant_op.constant([[10.0, 20.0], [-20.0, -10.0]])
     self._writer.histogram("image1", a)
     event = self._readLastEvent()
     self.assertEqual("image1", event.summary.value[0].tag)
     self.assertTrue(event.summary.value[0].image)

 def testHistogramSummary(self):
   with context.device(self._test_device):
     y = constant_op.constant([1.0, 3.0, 3.0, 7.0])
     self._writer.histogram("y", y)
     event = self._readLastEvent()
     self.assertEqual("y", event.summary.value[0].tag)
     self.assertTrue(event.summary.value[0].histo)

 def testScalarSummary(self):
   with context.device(self._test_device):
     x = constant_op.constant(1337.0)
     self._writer.scalar("x", x)
     event = self._readLastEvent()
     self.assertTrue("x", event.summary.value[0].tag)
     self.assertEqual(1337.0, event.summary.value[0].simple_value)

 def benchmark_defun_matmul_100_by_784_GPU(self):
   if not context.num_gpus():
     return
   with context.device(GPU):
     m = self._m_100_by_784.gpu()
     self._benchmark_defun_matmul(
         m, transpose_b=True, num_iters=self._num_iters_100_by_784)

 def fn(x):
   with context.device('/gpu:0'):
     b = constant_op.constant(2.0)
     c = math_ops.add(x.gpu(), b)
     # TODO(apassos): remove cpu below by making TensorVSPace aware
     # of devices.
     return math_ops.add(c, constant_op.constant(3.0)).cpu()

 def benchmark_read_variable_op_with_tape_2_by_2_GPU(self):
   if not context.num_gpus():
     return
   with context.device(GPU):
     m = resource_variable_ops.ResourceVariable(self._m_2_by_2.gpu())
     self._benchmark_read_variable_with_tape(
         m, num_iters=self._num_iters_2_by_2)

 def benchmark_defun_matmul_2_by_2_GPU(self):
   if not context.num_gpus():
     return
   with context.device(GPU):
     m = self._m_2_by_2.gpu()
     self._benchmark_defun_matmul(
         m, transpose_b=False, num_iters=self._num_iters_2_by_2)

  def __init__(self,
               logdir,
               max_queue=10,
               flush_secs=120,
               filename_suffix=""):
    """Summary writer for TensorBoard, compatible with eager execution.

    If necessary, multiple instances of `SummaryWriter` can be created, with
    distinct `logdir`s and `name`s. Each `SummaryWriter` instance will retain
    its independent `global_step` counter and data writing destination.

    Example:
    ```python
    writer = tfe.SummaryWriter("my_model")

    # ... Code that sets up the model and data batches ...

    for _ in xrange(train_iters):
      loss = model.train_batch(batch)
      writer.scalar("loss", loss)
      writer.step()
    ```

    Args:
      logdir: Directory in which summary files will be written.
      max_queue: Number of summary items to buffer before flushing to
        filesystem. If 0, summaries will be flushed immediately.
      flush_secs: Number of secondsbetween forced commits to disk.
      filename_suffix: Suffix of the event protobuf files in which the summary
        data are stored.

    Raises:
      ValueError: If this constructor is called not under eager execution.
    """
    # TODO(apassos, ashankar): Make this class and the underlying
    # contrib.summary_ops compatible with graph model and remove this check.
    if not context.in_eager_mode():
      raise ValueError(
          "Use of SummaryWriter is currently supported only with eager "
          "execution enabled. File an issue at "
          "https://github.com/tensorflow/tensorflow/issues/new to express "
          "interest in fixing this.")

    # TODO(cais): Consider adding name keyword argument, which if None or empty,
    # will register the global global_step that training_util.get_global_step()
    # can find.
    with context.device(self._CPU_DEVICE):
      self._name = uuid.uuid4().hex
      self._global_step = 0
      self._global_step_tensor = variable_scope.get_variable(
          "global_step/summary_writer/" + self._name,
          shape=[], dtype=dtypes.int64,
          initializer=init_ops.zeros_initializer())
      self._global_step_dirty = False
      self._resource = gen_summary_ops.summary_writer(shared_name=self._name)
      gen_summary_ops.create_summary_file_writer(
          self._resource, logdir, max_queue, flush_secs, filename_suffix)
      # Delete the resource when this object is deleted
      self._resource_deleter = resource_variable_ops.EagerResourceDeleter(
          handle=self._resource, handle_device=self._CPU_DEVICE)