Method testVirtualGpu

tensorflow/python/framework/config_test.py:443–484

(self)

  @test_util.run_gpu_only
  @reset_eager
  def testVirtualGpu(self):
    gpus = config.list_physical_devices('GPU')
    self.assertNotEqual(len(gpus), 0)

    self.assertIsNone(config.get_virtual_device_configuration(gpus[-1]))
    config.set_virtual_device_configuration(gpus[-1], [
        context.VirtualDeviceConfiguration(memory_limit=10),
        context.VirtualDeviceConfiguration(memory_limit=10)
    ])
    self.assertEqual(len(config.get_virtual_device_configuration(gpus[-1])), 2)

    logical_gpus = config.list_logical_devices('GPU')
    self.assertTrue(len(logical_gpus), len(gpus) + 1)
    for i in range(0, len(logical_gpus)):
      with ops.device('/device:GPU:' + str(i)):
        a = constant_op.constant(1.0)
        self.evaluate(a)

    with self.assertRaisesRegexp(RuntimeError, 'unknown device'):
      with ops.device('/device:GPU:' + str(len(logical_gpus))):
        a = constant_op.constant(1.0)
        self.evaluate(a)

    # Modifying the GPU configuration is not supported
    with self.assertRaisesRegexp(RuntimeError, 'cannot be modified'):
      config.set_virtual_device_configuration(gpus[-1], [
          context.VirtualDeviceConfiguration(memory_limit=20),
          context.VirtualDeviceConfiguration(memory_limit=20)
      ])

    with self.assertRaisesRegexp(RuntimeError, 'cannot be modified'):
      config.set_virtual_device_configuration(gpus[-1], [
          context.VirtualDeviceConfiguration(memory_limit=10),
          context.VirtualDeviceConfiguration(memory_limit=10),
          context.VirtualDeviceConfiguration(memory_limit=10)
      ])

    # Setting the same GPU configuration is fine
    config.set_virtual_device_configuration(gpus[-1], [
        context.VirtualDeviceConfiguration(memory_limit=10),
        context.VirtualDeviceConfiguration(memory_limit=10)
    ])

  @test_util.run_gpu_only
  @reset_eager