本文整理汇总了Python中tensorflow.python.distribute.cluster_resolver.TPUClusterResolver._get_device_dict_and_cores方法的典型用法代码示例。如果您正苦于以下问题：Python TPUClusterResolver._get_device_dict_and_cores方法的具体用法？Python TPUClusterResolver._get_device_dict_and_cores怎么用？Python TPUClusterResolver._get_device_dict_and_cores使用的例子？那么恭喜您, 这里精选的方法代码示例或许可以为您提供帮助。您也可以进一步了解该方法所在类tensorflow.python.distribute.cluster_resolver.TPUClusterResolver的用法示例。
在下文中一共展示了TPUClusterResolver._get_device_dict_and_cores方法的2个代码示例，这些例子默认根据受欢迎程度排序。您可以为喜欢或者感觉有用的代码点赞，您的评价将有助于系统推荐出更棒的Python代码示例。
示例1: testGetDeviceDictAndCoresWithCPUsAndGPUs
# 需要导入模块: from tensorflow.python.distribute.cluster_resolver import TPUClusterResolver [as 别名]
# 或者: from tensorflow.python.distribute.cluster_resolver.TPUClusterResolver import _get_device_dict_and_cores [as 别名]
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, {})示例2: testGetDeviceDictAndCoresWithTPUs
# 需要导入模块: from tensorflow.python.distribute.cluster_resolver import TPUClusterResolver [as 别名]
# 或者: from tensorflow.python.distribute.cluster_resolver.TPUClusterResolver import _get_device_dict_and_cores [as 别名]
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']})