From 904a8407865d69862ad993b7c7b67d530483851c Mon Sep 17 00:00:00 2001 From: Julien Chaumond Date: Fri, 5 Aug 2022 18:24:07 +0200 Subject: [PATCH 1/5] `pipeline` support for `device="mps"` (or any other string) --- src/transformers/pipelines/__init__.py | 6 +++++- src/transformers/pipelines/base.py | 12 ++++++++---- 2 files changed, 13 insertions(+), 5 deletions(-) diff --git a/src/transformers/pipelines/__init__.py b/src/transformers/pipelines/__init__.py index dfa75768d8f811..e9543dff320cd6 100755 --- a/src/transformers/pipelines/__init__.py +++ b/src/transformers/pipelines/__init__.py @@ -422,6 +422,7 @@ def pipeline( revision: Optional[str] = None, use_fast: bool = True, use_auth_token: Optional[Union[str, bool]] = None, + device: Optional[Union[int, str, "torch.device"]] = None, device_map=None, torch_dtype=None, trust_remote_code: Optional[bool] = None, @@ -508,6 +509,9 @@ def pipeline( use_auth_token (`str` or *bool*, *optional*): The token to use as HTTP bearer authorization for remote files. If `True`, will use the token generated when running `huggingface-cli login` (stored in `~/.huggingface`). + device (`int` or `str` or `torch.device`): + Sent directly as `model_kwargs` (just a simpler shortcut). Defines the device (*e.g.*, `"cpu"`, `"cuda:1"`, + `"mps"`, or a GPU ordinal rank like `1`) on which this pipeline will be allocated. device_map (`str` or `Dict[str, Union[int, str, torch.device]`, *optional*): Sent directly as `model_kwargs` (just a simpler shortcut). When `accelerate` library is present, set `device_map="auto"` to compute the most optimized `device_map` automatically. [More @@ -802,4 +806,4 @@ def pipeline( if feature_extractor is not None: kwargs["feature_extractor"] = feature_extractor - return pipeline_class(model=model, framework=framework, task=task, **kwargs) + return pipeline_class(model=model, framework=framework, task=task, device=device, **kwargs) diff --git a/src/transformers/pipelines/base.py b/src/transformers/pipelines/base.py index 6e2c28e5ddf84d..f3b5d83dab504e 100644 --- a/src/transformers/pipelines/base.py +++ b/src/transformers/pipelines/base.py @@ -704,7 +704,7 @@ def predict(self, X): Reference to the object in charge of parsing supplied pipeline parameters. device (`int`, *optional*, defaults to -1): Device ordinal for CPU/GPU supports. Setting this to -1 will leverage CPU, a positive will run the model on - the associated CUDA device id. You can pass native `torch.device` too. + the associated CUDA device id. You can pass native `torch.device` or a `str` too. binary_output (`bool`, *optional*, defaults to `False`): Flag indicating if the output the pipeline should happen in a binary format (i.e., pickle) or as raw text. """ @@ -747,7 +747,7 @@ def __init__( framework: Optional[str] = None, task: str = "", args_parser: ArgumentHandler = None, - device: int = -1, + device: Union[int, str, "torch.device"] = -1, binary_output: bool = False, **kwargs, ): @@ -763,11 +763,15 @@ def __init__( if is_torch_available() and isinstance(device, torch.device): self.device = device else: - self.device = device if framework == "tf" else torch.device("cpu" if device < 0 else f"cuda:{device}") + self.device = ( + device + if framework == "tf" + else torch.device(device if type(device) == str else "cpu" if device < 0 else f"cuda:{device}") + ) self.binary_output = binary_output # Special handling - if self.framework == "pt" and self.device.type == "cuda": + if self.framework == "pt" and self.device.type != "cpu": self.model = self.model.to(self.device) # Update config with task specific parameters From f0e3f047e776d9c6d859377e9f03410fbb76bf45 Mon Sep 17 00:00:00 2001 From: Julien Chaumond Date: Wed, 10 Aug 2022 18:20:56 +0200 Subject: [PATCH 2/5] Simplify `if` nesting --- src/transformers/pipelines/__init__.py | 4 ++-- src/transformers/pipelines/base.py | 17 ++++++++++------- 2 files changed, 12 insertions(+), 9 deletions(-) diff --git a/src/transformers/pipelines/__init__.py b/src/transformers/pipelines/__init__.py index e9543dff320cd6..b3d16a2e3aac5b 100755 --- a/src/transformers/pipelines/__init__.py +++ b/src/transformers/pipelines/__init__.py @@ -510,8 +510,8 @@ def pipeline( The token to use as HTTP bearer authorization for remote files. If `True`, will use the token generated when running `huggingface-cli login` (stored in `~/.huggingface`). device (`int` or `str` or `torch.device`): - Sent directly as `model_kwargs` (just a simpler shortcut). Defines the device (*e.g.*, `"cpu"`, `"cuda:1"`, - `"mps"`, or a GPU ordinal rank like `1`) on which this pipeline will be allocated. + Defines the device (*e.g.*, `"cpu"`, `"cuda:1"`, `"mps"`, or a GPU ordinal rank like `1`) on which this + pipeline will be allocated. device_map (`str` or `Dict[str, Union[int, str, torch.device]`, *optional*): Sent directly as `model_kwargs` (just a simpler shortcut). When `accelerate` library is present, set `device_map="auto"` to compute the most optimized `device_map` automatically. [More diff --git a/src/transformers/pipelines/base.py b/src/transformers/pipelines/base.py index f3b5d83dab504e..8d5de5cc9a6cf7 100644 --- a/src/transformers/pipelines/base.py +++ b/src/transformers/pipelines/base.py @@ -760,14 +760,17 @@ def __init__( self.feature_extractor = feature_extractor self.modelcard = modelcard self.framework = framework - if is_torch_available() and isinstance(device, torch.device): - self.device = device + if is_torch_available() and self.framework == "pt": + if isinstance(device, torch.device): + self.device = device + elif type(device) == str: + self.device = torch.device(device) + elif device < 0: + self.device = torch.device("cpu") + else: + self.device = torch.device("cuda:{device}") else: - self.device = ( - device - if framework == "tf" - else torch.device(device if type(device) == str else "cpu" if device < 0 else f"cuda:{device}") - ) + self.device = device self.binary_output = binary_output # Special handling From a7fe56eb469e6a56c88b4a662fdad7f5583c6eb6 Mon Sep 17 00:00:00 2001 From: Julien Chaumond Date: Wed, 10 Aug 2022 18:34:17 +0200 Subject: [PATCH 3/5] Update src/transformers/pipelines/base.py Co-authored-by: Sylvain Gugger <35901082+sgugger@users.noreply.github.com> --- src/transformers/pipelines/base.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/transformers/pipelines/base.py b/src/transformers/pipelines/base.py index 8d5de5cc9a6cf7..03ce3678c5b98f 100644 --- a/src/transformers/pipelines/base.py +++ b/src/transformers/pipelines/base.py @@ -763,7 +763,7 @@ def __init__( if is_torch_available() and self.framework == "pt": if isinstance(device, torch.device): self.device = device - elif type(device) == str: + elif isinstance(int, str): self.device = torch.device(device) elif device < 0: self.device = torch.device("cpu") From 4746db774407a79f8a2cbbfcd0c4019c92fc7e53 Mon Sep 17 00:00:00 2001 From: Julien Chaumond Date: Wed, 10 Aug 2022 18:38:10 +0200 Subject: [PATCH 4/5] Fix? @sgugger --- src/transformers/pipelines/base.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/transformers/pipelines/base.py b/src/transformers/pipelines/base.py index 03ce3678c5b98f..a0ce06ec5e33f1 100644 --- a/src/transformers/pipelines/base.py +++ b/src/transformers/pipelines/base.py @@ -763,7 +763,7 @@ def __init__( if is_torch_available() and self.framework == "pt": if isinstance(device, torch.device): self.device = device - elif isinstance(int, str): + elif isinstance(device, str): self.device = torch.device(device) elif device < 0: self.device = torch.device("cpu") From c23b7362f72a775fa0501c3e4b3ebc5b81c2a1dc Mon Sep 17 00:00:00 2001 From: Julien Chaumond Date: Wed, 10 Aug 2022 18:40:04 +0200 Subject: [PATCH 5/5] =?UTF-8?q?passing=20`attr=3DNone`=20is=20not=20the=20?= =?UTF-8?q?same=20as=20not=20passing=20`attr`=20=F0=9F=A4=AF?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- src/transformers/pipelines/__init__.py | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/src/transformers/pipelines/__init__.py b/src/transformers/pipelines/__init__.py index b3d16a2e3aac5b..80a48198f3fe6d 100755 --- a/src/transformers/pipelines/__init__.py +++ b/src/transformers/pipelines/__init__.py @@ -806,4 +806,7 @@ def pipeline( if feature_extractor is not None: kwargs["feature_extractor"] = feature_extractor - return pipeline_class(model=model, framework=framework, task=task, device=device, **kwargs) + if device is not None: + kwargs["device"] = device + + return pipeline_class(model=model, framework=framework, task=task, **kwargs)