diff --git a/algorithm_store.json b/algorithm_store.json index b8c7252..ec2057d 100644 --- a/algorithm_store.json +++ b/algorithm_store.json @@ -21,7 +21,7 @@ }, { "name": "partial", - "description": "Retrieves GPU accessibility details for TensorFlow.", + "description": "Retrieves GPU accessibility details for PyTorch.", "type": "federated", "databases": [], "arguments": [] diff --git a/docs/v6-gpu-accessibility/Implementation.rst b/docs/v6-gpu-accessibility/Implementation.rst index 21014d3..3d4f48d 100644 --- a/docs/v6-gpu-accessibility/Implementation.rst +++ b/docs/v6-gpu-accessibility/Implementation.rst @@ -20,9 +20,10 @@ node. ``partial`` ~~~~~~~~~~~ -The partial function retrieves GPU accessibility details for TensorFlow through the helper function. +The partial function retrieves GPU accessibility details for PyTorch through the helper function. -``get_tensorflow_gpu_details`` -~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ -This function checks if TensorFlow is accessible on the GPU. + +``get_pytorch_gpu_details`` +~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ +This function checks if PyTorch is accessible on the GPU. If a GPU is available it returns the GPU device name(s), memory, cores, and compute capability. diff --git a/docs/v6-gpu-accessibility/Validation.rst b/docs/v6-gpu-accessibility/Validation.rst index 4ac7754..44a781a 100644 --- a/docs/v6-gpu-accessibility/Validation.rst +++ b/docs/v6-gpu-accessibility/Validation.rst @@ -5,15 +5,16 @@ The following steps can be taken to verify the GPU accessibility: .. code-block:: python - import tensorflow as tf + import torch - gpus = tf.config.list_physical_devices('GPU') - if not gpus: - print('No GPUs available for TensorFlow') + + if not torch.cuda.is_available(): + print('No GPUs available for PyTorch') else: print({ - "number of GPUs": len(gpus), - "memory": f'{[tf.config.experimental.get_memory_info(gpu.name)['current'] for gpu in gpus]} bytes, - "cores": [tf.config.experimental.get_device_details(gpu)['core_count'] for gpu in gpus], - "compute capability": [tf.config.experimental.get_device_details(gpu)['compute_capability'] for gpu in gpus] + "number of GPUs": torch.cuda.device_count(), + "memory": f'{[torch.cuda.get_device_properties(i).total_memory for i in range(torch.cuda.device_count())]} bytes', + "cores": [torch.cuda.get_device_properties(i).multi_processor_count for i in range(torch.cuda.device_count())], + "compute capability": [torch.cuda.get_device_properties(i).major for i in range(torch.cuda.device_count())] }) + diff --git a/requirements.txt b/requirements.txt index 65ff5c9..2f61abd 100644 --- a/requirements.txt +++ b/requirements.txt @@ -1,3 +1,3 @@ vantage6-algorithm-tools pandas -tensorflow +torch diff --git a/setup.py b/setup.py index 36b70dc..23dad3c 100644 --- a/setup.py +++ b/setup.py @@ -20,6 +20,6 @@ install_requires=[ 'vantage6-algorithm-tools', 'pandas', - 'tensorflow' + 'torch' ] ) diff --git a/v6-gpu-accessibility/partial.py b/v6-gpu-accessibility/partial.py index 2444d91..d7a6814 100644 --- a/v6-gpu-accessibility/partial.py +++ b/v6-gpu-accessibility/partial.py @@ -6,7 +6,7 @@ encryption if that is enabled). From there, they are sent to the partial task or directly to the user (if they requested partial results). """ -import tensorflow as tf +import torch from typing import Any @@ -18,37 +18,36 @@ @algorithm_client def partial(client: AlgorithmClient) -> Any: """ - Retrieves GPU accessibility details for TensorFlow. + Retrieves GPU accessibility details for TensorFlow and PyTorch. Args: client (AlgorithmClient): The algorithm client instance. Returns: - dict: A dictionary containing the organisation ID and GPU accessibility details for TensorFlow. + dict: A dictionary containing the organisation ID and GPU accessibility details for TensorFlow and PyTorch. """ info("Retrieving details of any GPUs available to the algorithm's Docker container") return { 'organisation_id': client.organization_id, 'gpu_accessibility': { - "Tensorflow": get_tensorflow_gpu_details() + "PyTorch": get_pytorch_gpu_details() } } -def get_tensorflow_gpu_details() -> dict or str: +def get_pytorch_gpu_details() -> dict or str: """ - Retrieves details about available TensorFlow GPUs. + Retrieves details about available PyTorch GPUs. Returns: dict or str: A dictionary containing the number of GPUs, their names, memory, cores, and compute capability if GPUs are available. Otherwise, returns 'No GPUs available'. """ - gpus = tf.config.list_physical_devices('GPU') - if not gpus: + if not torch.cuda.is_available(): return 'No GPUs available' return { - "number of GPUs": len(gpus), - "memory": f"{[(tf.config.experimental.get_memory_info(gpu.name)['current']/ (1024 * 1024)) for gpu in gpus]} MB", - "cores": [tf.config.experimental.get_device_details(gpu)['core_count'] for gpu in gpus], - "compute capability": [tf.config.experimental.get_device_details(gpu)['compute_capability'] for gpu in gpus] + "number of GPUs": torch.cuda.device_count(), + "memory": f"{[(torch.cuda.get_device_properties(i).total_memory/ (1024 * 1024)) for i in range(torch.cuda.device_count())]} MB", + "cores": [torch.cuda.get_device_properties(i).multi_processor_count for i in range(torch.cuda.device_count())], + "compute capability": [torch.cuda.get_device_properties(i).major for i in range(torch.cuda.device_count())] }