Spaces:
Runtime error
Runtime error
File size: 1,221 Bytes
be11144 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 |
import tensorflow as tf
use_gpu = tf.test.is_gpu_available(
cuda_only=True,
min_cuda_compute_capability=None
)
cpu_device_id = 0
gpu_device_id = 0
def get_device_name():
"""
Get the current tensorflow device name we are using.
"""
global use_gpu
global cpu_device_id
global gpu_device_id
return '/device:gpu:' + str(gpu_device_id) if use_gpu else '/device:cpu:' + str(cpu_device_id)
def set_use_gpu(v: bool):
"""
Set whether to use CUDA or not.
"""
global use_gpu
use_gpu = v
def get_use_gpu():
"""
Get whether we are using CUDA or not.
"""
global use_gpu
return use_gpu
def set_cpu_device_id(did: int):
"""
Set the cpu device id we are using.
"""
global cpu_device_id
cpu_device_id = did
def get_cpu_device_id():
"""
Get the cpu device id we are using.
"""
global cpu_device_id
return cpu_device_id
def set_gpu_device_id(did: int):
"""
Set the gpu device id we are using.
"""
global gpu_device_id
gpu_device_id = did
def get_gpu_device_id():
"""
Get the gpu device id we are using.
"""
global gpu_device_id
return gpu_device_id
|