def __init__(self): default_device = '/' + os.path.basename(tf.zeros(()).device) default_device = ComputeDevice(self, default_device, default_device.split(":")[-2], -1, -1, "", default_device) for device in self.list_devices(): if device.name == default_device.name: default_device = device Backend.__init__(self, "TensorFlow", default_device)
def list_devices(self, device_type: str or None = None) -> List[ComputeDevice]: types = ['cpu', 'gpu', 'tpu'] if device_type is None else [device_type.lower()] devices = [] for device_type in types: try: for jax_dev in jax.devices(device_type): devices.append(ComputeDevice(self, jax_dev.device_kind, jax_dev.platform.upper(), -1, -1, f"id={jax_dev.id}", jax_dev)) except RuntimeError as err: pass # this is just Jax not finding anything. jaxlib.xla_client._get_local_backends() could help but isn't currently available on GitHub actions return devices
def list_devices(self, device_type: str or None = None) -> List[ComputeDevice]: devices = [] for jax_dev in jax.devices(): jax_dev_type = jax_dev.platform.upper() if device_type is None or device_type == jax_dev_type: description = f"id={jax_dev.id}" devices.append( ComputeDevice(self, jax_dev.device_kind, jax_dev_type, -1, -1, description, jax_dev)) return devices
def list_devices(self, device_type: str or None = None) -> List[ComputeDevice]: tf_devices = device_lib.list_local_devices() devices = [] for device in tf_devices: if device_type in (None, device.device_type): devices.append( ComputeDevice(self, device.name, device.device_type, device.memory_limit, processor_count=-1, description=str(device), ref=tf.device(device.name))) return devices
def list_devices(self, device_type: str or None = None) -> List[ComputeDevice]: devices = [] if device_type in (None, 'CPU'): devices.append(self.cpu) if device_type in (None, 'GPU'): for index in range(torch.cuda.device_count()): properties = torch.cuda.get_device_properties(index) devices.append(ComputeDevice(self, properties.name, 'GPU', properties.total_memory, properties.multi_processor_count, f"compute capability {properties.major}.{properties.minor}", ref=f'cuda:{index}')) return devices
def __init__(self): self.cpu = ComputeDevice(self, "CPU", 'CPU', -1, -1, "", ref='cpu') Backend.__init__(self, 'PyTorch', default_device=self.cpu)
def __init__(self): cpu = NUMPY.cpu self.cpu = ComputeDevice(self, "CPU", 'CPU', cpu.memory, cpu.processor_count, cpu.description, ref='cpu') gpus = self.list_devices('GPU') Backend.__init__(self, 'PyTorch', default_device=gpus[0] if gpus else cpu)