refactor: unify cpu gpu into device

CHANGELOG.md

@@ -16,6 +16,8 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ### Changed
 
+- Use `device` parameter to replace `cpu` to align with docarray  ([#577](https://github.com/jina-ai/finetuner/pull/577))
+
 ### Fixed
 
 ### Docs

docs/tasks/image-to-image.md

@@ -57,7 +57,7 @@ run = finetuner.fit(
     batch_size=128,
     epochs=5,
     learning_rate=1e-5,
-    cpu=False,
+    device='cuda',
     callbacks=[
         EvaluationCallback(
             query_data='tll-test-query-da',

docs/tasks/text-to-image.md

@@ -45,7 +45,7 @@ run = finetuner.fit(
     epochs=5,
     learning_rate= 1e-7,
     loss='CLIPLoss',
-    cpu=False,
+    device='cuda',
 )
 ```
 Let's understand what this piece of code does:

docs/tasks/text-to-text.md

@@ -115,7 +115,7 @@ run = finetuner.fit(
     learning_rate = 1e-5,
     epochs=3,
     batch_size=128,
-    cpu=False,
+    device='cuda',
     callbacks=[
         EvaluationCallback(
             query_data='quora_query_dev.da',

docs/walkthrough/run-job.md

@@ -71,7 +71,7 @@ run = finetuner.fit(
     scheduler_step='batch',
     freeze=False, # If applied will freeze the embedding model, only train the MLP.
     output_dim=512, # Attach a MLP on top of embedding model.
-    cpu=False,
+    device='cuda',
     num_workers=4,
     to_onnx=False,  # If set, please pass `is_onnx` when making inference.
 )

docs/walkthrough/using-callbacks.md

@@ -108,7 +108,7 @@ run = finetuner.fit(
     epochs=10,
     learning_rate= 1e-5,
     loss='CLIPLoss',
-    cpu=False,
+    device='cuda',
     callbacks= [
         callback.EarlyStopping(
             monitor = "train_loss",

finetuner/__init__.py

@@ -110,7 +110,8 @@ def fit(
     scheduler_step: str = 'batch',
     freeze: bool = False,
     output_dim: Optional[int] = None,
-    cpu: bool = True,
+    cpu: bool = False,
+    device: str = 'cuda',
     num_workers: int = 4,
     to_onnx: bool = False,
 ) -> Run:
@@ -167,10 +168,16 @@ def fit(
     :param output_dim: The expected output dimension as `int`.
         If set, will attach a projection head.
     :param cpu: Whether to use the CPU. If set to `False` a GPU will be used.
+    :param device: Whether to use the CPU, if set to `cuda`, a Nvidia GPU will be used.
+        otherwise use `cpu` to run a cpu job.
     :param num_workers: Number of CPU workers. If `cpu: False` this is the number of
         workers used by the dataloader.
     :param to_onnx: If the model is an onnx model or not. If you call the `fit` function
         with `to_onnx=True`, please set this parameter as `True`.
+
+    .. note::
+       Unless necessary, please stick with `device="cuda"`, `cpu` training could be
+       extremely slow and inefficient.
     """
     return ft.create_run(
         model=model,
@@ -193,6 +200,7 @@ def fit(
         freeze=freeze,
         output_dim=output_dim,
         cpu=cpu,
+        device=device,
         num_workers=num_workers,
         to_onnx=to_onnx,
     )
@@ -304,7 +312,7 @@ def get_model(
     token: Optional[str] = None,
     batch_size: int = 32,
     select_model: Optional[str] = None,
-    gpu: bool = False,
+    device: str = 'cpu',
     logging_level: str = 'WARNING',
     is_onnx: bool = False,
 ):
@@ -323,7 +331,8 @@ def get_model(
     :param select_model: Finetuner run artifacts might contain multiple models. In
         such cases you can select which model to deploy using this argument. For CLIP
         fine-tuning, you can choose either `clip-vision` or `clip-text`.
-    :param gpu: if specified to True, use cuda device for inference.
+    :param device: Whether to use the CPU, if set to `cuda`, a Nvidia GPU will be used.
+        otherwise use `cpu` to run a cpu job.
     :param logging_level: The executor logging level. See
         https://docs.python.org/3/library/logging.html#logging-levels for available
         options.
@@ -338,7 +347,7 @@ def get_model(
         TorchInferenceEngine,
     )
 
-    if gpu:
+    if device == 'cuda' and is_onnx:
         warnings.warn(
             message='You are using cuda device for ONNX inference, please consider'
             'call `pip install onnxruntime-gpu` to speed up inference.',
@@ -350,7 +359,7 @@ def get_model(
             token=token,
             batch_size=batch_size,
             select_model=select_model,
-            device='cuda' if gpu else 'cpu',
+            device=device,
             logging_level=logging_level,
         )
     else:
@@ -359,7 +368,7 @@ def get_model(
             token=token,
             batch_size=batch_size,
             select_model=select_model,
-            device='cuda' if gpu else 'cpu',
+            device=device,
             logging_level=logging_level,
         )
     return inference_engine

finetuner/experiment.py

@@ -1,3 +1,4 @@
+import warnings
 from dataclasses import fields
 from typing import Any, Dict, List, Optional, Union
 
@@ -13,6 +14,7 @@
     CREATED_AT,
     DATA,
     DESCRIPTION,
+    DEVICE,
     EPOCHS,
     EVAL_DATA,
     EXPERIMENT_NAME,
@@ -165,14 +167,24 @@ def create_run(
             **kwargs,
         )
 
-        cpu = kwargs.get(CPU, True)
+        device = kwargs.get(DEVICE, 'cuda')
+        if device == 'cuda' and kwargs.get(CPU, True):
+            device = 'gpu'
+            warnings.warn(
+                message='Parameter `cpu=True` will be deprecated from Finetuner 0.7.0,'
+                'please use `device="cpu" or `device="cuda" instead.`',
+                category=DeprecationWarning,
+            )
+        if device == 'cuda':
+            device = 'gpu'  # Map cuda to gpu to align with api
+
         num_workers = kwargs.get(NUM_WORKERS, 4)
 
         run_info = self._client.create_run(
             run_name=run_name,
             experiment_name=self._name,
             run_config=config,
-            device='cpu' if cpu else 'gpu',
+            device=device,
             cpus=num_workers,
             gpus=1,
         )

finetuner/finetuner.py

@@ -170,7 +170,8 @@ def create_run(
         scheduler_step: str = 'batch',
         freeze: bool = False,
         output_dim: Optional[int] = None,
-        cpu: bool = True,
+        cpu: bool = False,
+        device: str = 'cuda',
         num_workers: int = 4,
         to_onnx: bool = False,
     ) -> Run:
@@ -205,6 +206,7 @@ def create_run(
             freeze=freeze,
             output_dim=output_dim,
             cpu=cpu,
+            device=device,
             num_workers=num_workers,
             to_onnx=to_onnx,
         )

tests/unit/test_experiment.py

@@ -148,6 +148,7 @@ def test_create_run_config():
         image_modality=None,
         text_modality=None,
         cpu=False,
+        device='cuda',
         wandb_api_key=None,
     )
     assert config == expected_config