Merge pull request #1185 from rajveer43/docs-update

Improve documentation of device-related methods and docs regarding GPU usage

lightwood/helpers/device.py

@@ -5,6 +5,24 @@
 
 
 def is_cuda_compatible():
+    """
+    Check if the system has CUDA-compatible devices with the required architecture and
+    compiled CUDA version.
+
+    This function checks the compatibility of CUDA devices available on the system by
+    comparing their architectures and the compiled CUDA version. It iterates through
+    the available devices and verifies if their architectures meet the minimum
+    requirement specified by the function, and also checks if the compiled CUDA
+    version is greater than a specific version.
+
+    Returns:
+        bool: True if there are compatible CUDA devices, otherwise False.
+
+    Example:
+        >>> is_compatible = is_cuda_compatible()
+        >>> print(is_compatible)
+        True
+    """
     compatible_device_count = 0
     if torch.version.cuda is not None:
         for d in range(device_count()):
@@ -23,6 +41,27 @@ def is_cuda_compatible():
 
 
 def get_devices():
+    """
+    Get the appropriate Torch device(s) based on CUDA availability and compatibility.
+
+    This function determines the appropriate Torch device(s) to be used for
+    computations based on the availability of CUDA and compatible devices. It checks
+    if CUDA is available and if the available CUDA devices are compatible according to
+    the 'is_cuda_compatible()' function. If compatible devices are found, the function
+    selects either the first available CUDA device or a randomly selected one based on
+    the 'RANDOM_GPU' environment variable. If CUDA is not available or no compatible
+    devices are found, the function returns the CPU device.
+
+    Returns:
+        Tuple: A tuple containing the selected Torch device and the number of available
+        devices.
+    Example:
+        >>> device, num_devices = get_devices()
+        >>> print(device)
+        cuda:0
+        >>> print(num_devices)
+        1
+    """
     if torch.cuda.is_available() and is_cuda_compatible():
         device_str = "cuda"
         available_devices = torch.cuda.device_count()
@@ -40,10 +79,22 @@ def get_devices():
 
 def get_device_from_name(device_name=''):
     """
-    Returns the device specified as an argument.
-    If the argument is left empty it will returns the output of get_devices().
-    
-    :param device_name: name of the device to use (default is an empty string), if is an empty string will use the output of get_devices() instead")
+    Get a Torch device based on the specified device name or default behavior.
+
+    This function returns a Torch device based on the specified device name or the
+    default behavior, which is to return the output of the 'get_devices()' function.
+
+    Args:
+        device_name (str, optional): Name of the device to use. Default is an empty
+        string.
+        
+    Returns:
+        torch.device: The selected Torch device.
+
+    Example:
+        >>> device = get_device_from_name('cuda:1')
+        >>> print(device)
+        cuda:1
     """ # noqa E501
     if(device_name != ''):
         device = torch.device(device_name)

lightwood/helpers/torch.py

@@ -5,6 +5,32 @@
 
 
 def concat_vectors_and_pad(vec_list, max_):
+    """
+    Concatenates a list of input vectors and pads them to match a specified maximum
+    length.
+
+    This function takes a list of input vectors, concatenates them along a specified
+    dimension (dim=0), and then pads the concatenated vector to achieve a specified
+    maximum length. The padding is done with zeros.
+
+    Args:
+        vec_list (list of torch.Tensor): List of input vectors to concatenate and pad.
+        max_ (int): The maximum length of the concatenated and padded vector.
+
+    Returns:
+        torch.Tensor: The concatenated and padded vector.
+
+    Raises:
+        AssertionError: If the length of 'vec_list' is not greater than 0, or if it
+                        exceeds 'max_len', or if 'max_len' is not greater than 0.
+
+    Example:
+        >>> input_tensors = [torch.tensor([1, 2]), torch.tensor([3, 4, 5])]
+        >>> max_length = 5
+        >>> concatenated_padded = concat_vectors_and_pad(input_tensors, max_length)
+        >>> print(concatenated_padded)
+        tensor([1, 2, 3, 4, 5])
+    """
     assert len(vec_list) > 0
     assert len(vec_list) <= max_
     assert max_ > 0
@@ -27,10 +53,29 @@ class LightwoodAutocast:
     """
     Equivalent to torch.cuda.amp.autocast, but checks device compute capability
     to activate the feature only when the GPU has tensor cores to leverage AMP.
+
+    **Attributes:**
+
+    * `active` (bool): Whether AMP is currently active. This attribute is at the class
+    level
+
+    **Usage:**
+
+    ```python
+    >>> import lightwood.helpers.torch as lt
+    >>> with lt.LightwoodAutocast():
+    ...     # This code will be executed in AMP mode.
+    ...     pass
     """
     active = False
 
     def __init__(self, enabled=True):
+        """
+        Initializes the context manager for Automatic Mixed Precision (AMP) functionality.
+
+        Args:
+            enabled (bool, optional): Whether to enable AMP. Defaults to True.
+        """
         self.major = 0  # GPU major version
         torch_version = [int(i) for i in torch.__version__.split('.')[:-1]]
 
@@ -50,12 +95,18 @@ def __init__(self, enabled=True):
         LightwoodAutocast.active = self._enabled
 
     def __enter__(self):
+        """
+        * `__enter__()`: Enters the context manager and enables AMP if it is not already enabled.
+        """
         if self._enabled:
             self.prev = torch.is_autocast_enabled()
             torch.set_autocast_enabled(self._enabled)
             torch.autocast_increment_nesting()
 
     def __exit__(self, *args):
+        """
+        * `__exit__()`: Exits the context manager and disables AMP.
+        """
         if self._enabled:
             # Drop the cache when we exit to a nesting level that's outside any instance of autocast
             if torch.autocast_decrement_nesting() == 0:
@@ -64,6 +115,9 @@ def __exit__(self, *args):
         return False
 
     def __call__(self, func):
+        """
+        * `__call__(self, func)`: Returns a decorated function that enables AMP when it is called.
+        """
         @functools.wraps(func)
         def decorate_autocast(*args, **kwargs):
             with self: