experiments/STT-tensorflow
1
0
Fork 0
Code Issues Pull Requests Projects Releases Wiki Activity

Switch tf.config API examples to use doctest

Browse Source 
PiperOrigin-RevId: 269682976
This commit is contained in:
Gaurav Jain 2019-09-17 16:40:16 -07:00 committed by TensorFlower Gardener
parent c52ee536c0
commit d44f3a37ec
1 changed files with 98 additions and 101 deletions

199
tensorflow/python/framework/config.py
View File

@ -304,10 +304,9 @@ def list_physical_devices(device_type=None):
The following example ensures the machine can see at least 1 GPU.
```python
physical_devices = tf.config.experimental.list_physical_devices('GPU')
assert len(physical_devices) > 0, "No GPUs found."
```
>>> physical_devices = tf.config.experimental.list_physical_devices('GPU')
>>> print("Num GPUs:", len(physical_devices))
Num GPUs: ...
Args:
device_type: (optional) Device type to filter by such as "CPU" or "GPU"
@ -328,15 +327,14 @@ def list_logical_devices(device_type=None):
For example:
```python
logical_devices = tf.config.experimental.list_logical_devices('GPU')
# Allocate on GPU:0
with tf.device(logical_devices[0].name):
one = tf.constant(1)
# Allocate on GPU:1
with tf.device(logical_devices[1].name):
two = tf.constant(2)
```
>>> logical_devices = tf.config.experimental.list_logical_devices('GPU')
>>> if len(logical_devices) > 0:
... # Allocate on GPU:0
... with tf.device(logical_devices[0].name):
... one = tf.constant(1)
... # Allocate on GPU:1
... with tf.device(logical_devices[1].name):
... two = tf.constant(2)
Args:
device_type: (optional) Device type to filter by such as "CPU" or "GPU"
@ -357,15 +355,16 @@ def get_visible_devices(device_type=None):
The following example verifies all visible GPUs have been disabled:
```python
physical_devices = config.experimental.list_physical_devices('GPU')
assert len(physical_devices) > 0, "Not enough GPU hardware devices available"
# Disable all GPUS
tf.config.experimental.set_visible_devices([], 'GPU')
visible_devices = tf.config.experimental.get_visible_devices()
for device in visible_devices:
assert device.device_type != 'GPU'
```
>>> physical_devices = tf.config.experimental.list_physical_devices('GPU')
>>> try:
... # Disable all GPUS
... tf.config.experimental.set_visible_devices([], 'GPU')
... visible_devices = tf.config.experimental.get_visible_devices()
... for device in visible_devices:
... assert device.device_type != 'GPU'
... except:
... # GPU not available or cannot modify virtual devices once initialized.
... pass
Args:
device_type: (optional) Device types to limit query to.
@ -388,15 +387,16 @@ def set_visible_devices(devices, device_type=None):
The following example demonstrates disabling the first GPU on the machine.
```python
physical_devices = config.experimental.list_physical_devices('GPU')
assert len(physical_devices) > 0, "Not enough GPU hardware devices available"
# Disable first GPU
tf.config.experimental.set_visible_devices(physical_devices[1:], 'GPU')
logical_devices = config.experimental.list_logical_devices('GPU')
# Logical device was not created for first GPU
assert len(logical_devices) == len(physical_devices) - 1
```
>>> physical_devices = tf.config.experimental.list_physical_devices('GPU')
>>> try:
... # Disable first GPU
... tf.config.experimental.set_visible_devices(physical_devices[1:], 'GPU')
... logical_devices = tf.config.experimental.list_logical_devices('GPU')
... # Logical device was not created for first GPU
... assert len(logical_devices) == len(physical_devices) - 1
... except:
... # GPU not available or cannot modify virtual devices once initialized.
... pass
Args:
devices: (optional) List of PhysicalDevice objects to make visible
@ -415,12 +415,13 @@ def get_memory_growth(device):
For example:
```python
physical_devices = config.experimental.list_physical_devices('GPU')
assert len(physical_devices) > 0, "Not enough GPU hardware devices available"
tf.config.experimental.set_memory_growth(physical_devices[0], True)
assert tf.config.experimental.get_memory_growth(physical_devices[0]) == True
```
>>> physical_devices = tf.config.experimental.list_physical_devices('GPU')
>>> try:
... tf.config.experimental.set_memory_growth(physical_devices[0], True)
... assert tf.config.experimental.get_memory_growth(physical_devices[0])
... except:
... # GPU not available or cannot modify virtual devices once initialized.
... pass
Args:
device: PhysicalDevice to query
@ -441,11 +442,12 @@ def set_memory_growth(device, enable):
For example:
```python
physical_devices = tf.config.experimental.list_physical_devices('GPU')
assert len(physical_devices) > 0, "Not enough GPU hardware devices available"
tf.config.experimental.set_memory_growth(physical_devices[0], True)
```
>>> physical_devices = tf.config.experimental.list_physical_devices('GPU')
>>> try:
... tf.config.experimental.set_memory_growth(physical_devices[0], True)
... except:
... # GPU not available or cannot modify virtual devices once initialized.
... pass
Args:
device: PhysicalDevice to configure
@ -463,20 +465,22 @@ def get_virtual_device_configuration(device):
For example:
```python
physical_devices = tf.config.experimental.list_physical_devices('CPU')
assert len(physical_devices) == 1, "No CPUs found"
configs = tf.config.experimental.get_virtual_device_configuration(
physical_devices[0])
assert configs is None
tf.config.experimental.set_virtual_device_configuration(
physical_devices[0],
[tf.config.experimental.VirtualDeviceConfiguration(),
tf.config.experimental.VirtualDeviceConfiguration()])
configs = tf.config.experimental.get_virtual_device_configuration(
physical_devices[0])
assert len(configs) == 2
```
>>> physical_devices = tf.config.experimental.list_physical_devices('CPU')
>>> assert len(physical_devices) == 1, "No CPUs found"
>>> configs = tf.config.experimental.get_virtual_device_configuration(
... physical_devices[0])
>>> try:
... assert configs is None
... tf.config.experimental.set_virtual_device_configuration(
... physical_devices[0],
... [tf.config.experimental.VirtualDeviceConfiguration(),
... tf.config.experimental.VirtualDeviceConfiguration()])
... configs = tf.config.experimental.get_virtual_device_configuration(
... physical_devices[0])
... assert len(configs) == 2
... except:
... # Cannot modify virtual devices once initialized.
... pass
Args:
device: PhysicalDevice to query
@ -500,54 +504,47 @@ def set_virtual_device_configuration(device, virtual_devices):
The following example splits the CPU into 2 virtual devices:
```python
physical_devices = tf.config.experimental.list_physical_devices('CPU')
assert len(physical_devices) == 1, "No CPUs found"
# Specify 2 virtual CPUs. Note currently memory limit is not supported.
tf.config.experimental.set_virtual_device_configuration(
physical_devices[0],
[tf.config.experimental.VirtualDeviceConfiguration(),
tf.config.experimental.VirtualDeviceConfiguration()])
logical_devices = tf.config.experimental.list_logical_devices('CPU')
assert len(logical_devices) == 2
try:
tf.config.experimental.set_virtual_device_configuration(
physical_devices[0],
[tf.config.experimental.VirtualDeviceConfiguration(),
tf.config.experimental.VirtualDeviceConfiguration(),
tf.config.experimental.VirtualDeviceConfiguration(),
tf.config.experimental.VirtualDeviceConfiguration()])
except:
print('Cannot modify the virtual devices once they have been initialized.')
```
>>> physical_devices = tf.config.experimental.list_physical_devices('CPU')
>>> assert len(physical_devices) == 1, "No CPUs found"
>>> # Specify 2 virtual CPUs. Note currently memory limit is not supported.
>>> try:
... tf.config.experimental.set_virtual_device_configuration(
... physical_devices[0],
... [tf.config.experimental.VirtualDeviceConfiguration(),
... tf.config.experimental.VirtualDeviceConfiguration()])
... logical_devices = tf.config.experimental.list_logical_devices('CPU')
... assert len(logical_devices) == 2
...
... tf.config.experimental.set_virtual_device_configuration(
... physical_devices[0],
... [tf.config.experimental.VirtualDeviceConfiguration(),
... tf.config.experimental.VirtualDeviceConfiguration(),
... tf.config.experimental.VirtualDeviceConfiguration(),
... tf.config.experimental.VirtualDeviceConfiguration()])
... except:
... # Cannot modify virtual devices once initialized.
... pass
The following example splits the GPU into 2 virtual devices with 100 MB each:
```python
physical_devices = tf.config.experimental.list_physical_devices('GPU')
assert len(physical_devices) > 0, "No GPUs found"
tf.config.experimental.set_virtual_device_configuration(
physical_devices[0],
[tf.config.experimental.VirtualDeviceConfiguration(memory_limit=100),
tf.config.experimental.VirtualDeviceConfiguration(memory_limit=100)])
try:
tf.config.experimental.set_memory_growth(physical_devices[0], True)
except:
print('Cannot set memory growth when virtual devices configured')
logical_devices = tf.config.experimental.list_logical_devices('GPU')
assert len(logical_devices) == len(physical_devices) + 1
try:
tf.config.experimental.set_virtual_device_configuration(
physical_devices[0],
[tf.config.experimental.VirtualDeviceConfiguration(memory_limit=10),
tf.config.experimental.VirtualDeviceConfiguration(memory_limit=10)])
except:
print('Cannot modify the virtual devices once they have been initialized.')
```
>>> physical_devices = tf.config.experimental.list_physical_devices('GPU')
>>> try:
... tf.config.experimental.set_virtual_device_configuration(
... physical_devices[0],
... [tf.config.experimental.VirtualDeviceConfiguration(memory_limit=100),
... tf.config.experimental.VirtualDeviceConfiguration(memory_limit=100)])
...
... tf.config.experimental.set_memory_growth(physical_devices[0], True)
... logical_devices = tf.config.experimental.list_logical_devices('GPU')
... assert len(logical_devices) == len(physical_devices) + 1
...
... tf.config.experimental.set_virtual_device_configuration(
... physical_devices[0],
... [tf.config.experimental.VirtualDeviceConfiguration(memory_limit=10),
... tf.config.experimental.VirtualDeviceConfiguration(memory_limit=10)])
... except:
... # GPU not available or cannot modify virtual devices once initialized.
... pass
Args:
device: (optional) Need to update