Python TPUClusterResolver._get_device_dict_and_cores примеры использования

Пример #1

    def testGetDeviceDictAndCoresWithTPUs(self):
        device_names = [
            '/job:tpu_worker/task:0/device:TPU:0',
            '/job:tpu_worker/task:1/device:TPU:1',
            '/job:tpu_worker/task:2/device:TPU:0',
            '/job:tpu_worker/task:3/device:TPU:1',
            '/job:tpu_worker/task:0/device:TPU:4',
            '/job:tpu_worker/task:1/device:TPU:5',
            '/job:tpu_worker/task:2/device:TPU:4',
            '/job:tpu_worker/task:3/device:TPU:5',
        ]
        device_list = [
            session._DeviceAttributes(name, 'TPU', 1024, 0)
            for name in device_names
        ]

        device_details = TPUClusterResolver._get_device_dict_and_cores(
            device_list)
        self.assertEqual(device_details.total_cores, 8)
        self.assertEqual(device_details.device_map, {
            '0': ['0', '4'],
            '1': ['1', '5'],
            '2': ['0', '4'],
            '3': ['1', '5']
        })

Пример #2

    def testGetDeviceDictAndCoresWithCPUsAndGPUs(self):
        device_names = [
            '/job:tpu_worker/task:0/device:CPU:0',
            '/job:tpu_worker/task:1/device:CPU:0',
            '/job:tpu_worker/task:2/device:CPU:0',
            '/job:tpu_worker/task:3/device:CPU:0',
            '/job:tpu_worker/task:0/device:GPU:1',
            '/job:tpu_worker/task:1/device:GPU:1',
            '/job:tpu_worker/task:2/device:GPU:1',
            '/job:tpu_worker/task:3/device:GPU:1',
        ]
        device_list = [
            session._DeviceAttributes(name, 'XLA', 1024, 0)
            for name in device_names
        ]

        device_dict, num_cores = TPUClusterResolver._get_device_dict_and_cores(
            device_list)
        self.assertEqual(num_cores, 0)
        self.assertEqual(device_dict, {})

Пример #3

Файл: tpu_cluster_resolver_test.py Проект: terrytangyuan/tensorflow

  def testGetDeviceDictAndCoresWithCPUsAndGPUs(self):
    device_names = [
        '/job:tpu_worker/task:0/device:CPU:0',
        '/job:tpu_worker/task:1/device:CPU:0',
        '/job:tpu_worker/task:2/device:CPU:0',
        '/job:tpu_worker/task:3/device:CPU:0',
        '/job:tpu_worker/task:0/device:GPU:1',
        '/job:tpu_worker/task:1/device:GPU:1',
        '/job:tpu_worker/task:2/device:GPU:1',
        '/job:tpu_worker/task:3/device:GPU:1',
    ]
    device_list = [
        session._DeviceAttributes(
            name, 'XLA', 1024, 0) for name in device_names
    ]

    device_dict, num_cores = TPUClusterResolver._get_device_dict_and_cores(
        device_list)
    self.assertEqual(num_cores, 0)
    self.assertEqual(device_dict, {})

Пример #4

Файл: tpu_cluster_resolver_test.py Проект: terrytangyuan/tensorflow

  def testGetDeviceDictAndCoresWithTPUs(self):
    device_names = [
        '/job:tpu_worker/task:0/device:TPU:0',
        '/job:tpu_worker/task:1/device:TPU:1',
        '/job:tpu_worker/task:2/device:TPU:0',
        '/job:tpu_worker/task:3/device:TPU:1',
        '/job:tpu_worker/task:0/device:TPU:4',
        '/job:tpu_worker/task:1/device:TPU:5',
        '/job:tpu_worker/task:2/device:TPU:4',
        '/job:tpu_worker/task:3/device:TPU:5',
    ]
    device_list = [
        session._DeviceAttributes(
            name, 'TPU', 1024, 0) for name in device_names
    ]

    device_details = TPUClusterResolver._get_device_dict_and_cores(
        device_list)
    self.assertEqual(device_details.total_cores, 8)
    self.assertEqual(device_details.device_map,
                     {'0': ['0', '4'],
                      '1': ['1', '5'],
                      '2': ['0', '4'],
                      '3': ['1', '5']})