From a2178d00b28052684e14c8fc4c0b3818a47b3ba7 Mon Sep 17 00:00:00 2001 From: Samhita Alla Date: Wed, 15 Feb 2023 15:30:01 +0530 Subject: [PATCH] Fix PyTorch transformer Signed-off-by: Samhita Alla --- flytekit/extras/pytorch/native.py | 32 ++++++++++++------- .../extras/pytorch/test_transformations.py | 1 + 2 files changed, 22 insertions(+), 11 deletions(-) diff --git a/flytekit/extras/pytorch/native.py b/flytekit/extras/pytorch/native.py index 4cf37871fb..cbaa0c80f0 100644 --- a/flytekit/extras/pytorch/native.py +++ b/flytekit/extras/pytorch/native.py @@ -1,5 +1,5 @@ import pathlib -from typing import Generic, Type, TypeVar +from typing import Type, TypeVar import torch @@ -12,7 +12,7 @@ T = TypeVar("T") -class PyTorchTypeTransformer(TypeTransformer, Generic[T]): +class PyTorchTypeTransformer(TypeTransformer[T]): def get_literal_type(self, t: Type[T]) -> LiteralType: return LiteralType( blob=_core_types.BlobType( @@ -63,30 +63,40 @@ def to_python_value(self, ctx: FlyteContext, lv: Literal, expected_python_type: # load pytorch tensor/module from a file return torch.load(local_path, map_location=map_location) - def guess_python_type(self, literal_type: LiteralType) -> Type[T]: + +class PyTorchTensorTransformer(PyTorchTypeTransformer[torch.Tensor]): + PYTORCH_FORMAT = "PyTorchTensor" + + def __init__(self): + super().__init__(name="PyTorch Tensor", t=torch.Tensor) + + def guess_python_type(self, literal_type: LiteralType) -> Type[torch.Tensor]: if ( literal_type.blob is not None and literal_type.blob.dimensionality == _core_types.BlobType.BlobDimensionality.SINGLE and literal_type.blob.format == self.PYTORCH_FORMAT ): - return T + return torch.Tensor raise ValueError(f"Transformer {self} cannot reverse {literal_type}") -class PyTorchTensorTransformer(PyTorchTypeTransformer[torch.Tensor]): - PYTORCH_FORMAT = "PyTorchTensor" - - def __init__(self): - super().__init__(name="PyTorch Tensor", t=torch.Tensor) - - class PyTorchModuleTransformer(PyTorchTypeTransformer[torch.nn.Module]): PYTORCH_FORMAT = "PyTorchModule" def __init__(self): super().__init__(name="PyTorch Module", t=torch.nn.Module) + def guess_python_type(self, literal_type: LiteralType) -> Type[torch.nn.Module]: + if ( + literal_type.blob is not None + and literal_type.blob.dimensionality == _core_types.BlobType.BlobDimensionality.SINGLE + and literal_type.blob.format == self.PYTORCH_FORMAT + ): + return torch.nn.Module + + raise ValueError(f"Transformer {self} cannot reverse {literal_type}") + TypeEngine.register(PyTorchTensorTransformer()) TypeEngine.register(PyTorchModuleTransformer()) diff --git a/tests/flytekit/unit/extras/pytorch/test_transformations.py b/tests/flytekit/unit/extras/pytorch/test_transformations.py index 9724a01182..a470b646d4 100644 --- a/tests/flytekit/unit/extras/pytorch/test_transformations.py +++ b/tests/flytekit/unit/extras/pytorch/test_transformations.py @@ -40,6 +40,7 @@ def test_get_literal_type(transformer, python_type, format): tf = transformer lt = tf.get_literal_type(python_type) assert lt == LiteralType(blob=BlobType(format=format, dimensionality=BlobType.BlobDimensionality.SINGLE)) + assert tf.guess_python_type(lt) == python_type @pytest.mark.parametrize(