PyTorch release

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": []

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.

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())]
               })
+

requirements.txt

@@ -1,3 +1,3 @@
 vantage6-algorithm-tools
 pandas
-tensorflow
+torch

setup.py

@@ -20,6 +20,6 @@
     install_requires=[
         'vantage6-algorithm-tools',
         'pandas',
-        'tensorflow'
+        'torch'
     ]
 )

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())]
     }