palantir_models.serializers 참조

제공되는 직렬화기는 @auto_serialize에 지정하여 모델 생성자 입력을 저장하고 복원합니다. 같은 형식의 모델을 여러 팀에서 반복 사용한다면 공통 직렬화기를 공유 라이브러리로 게시할 수 있습니다.

ModelSerializer 구현 계약

palantir_models.models._serialization.ModelSerializer를 확장합니다. 생성자는 사용자가 지정할 형식과 옵션을 받습니다. serialize(writer, obj)는 실제 객체를 ModelStateWriter에 기록하고 deserialize(reader)는 ModelStateReader에서 저장 상태를 읽어 복원한 객체를 반환합니다.

Writer의 open으로 파일을 쓰고 mkdir로 디렉터리를 만들 수 있습니다. Reader는 open으로 파일을 읽고 dir로 저장 디렉터리를 가져옵니다. 구체적인 상태 저장 및 로드 계약을 함께 사용합니다. 아래 구현은 각 직렬화기가 저장하는 파일과 전달하는 옵션을 보여 줍니다. 일반 모델에서는 라이브러리에 제공되는 클래스를 가져와 사용합니다.

CloudPickleSerializer

생성할 때 cloudpickle을 가져옵니다. cloudpickle.pkl을 바이너리 모드로 쓰고 읽으며 dump와 load에 같은 파일을 전달합니다.

import importlib
from types import ModuleType
from palantir_models.models._serialization import ModelSerializer
from palantir_models.models._state_accessors import ModelStateWriter, ModelStateReader


class CloudPickleSerializer(ModelSerializer[object]):
    """Serializer utilizing the cloudpickle library for generic objects"""

    file_name = "cloudpickle.pkl"
    cloudpickle: ModuleType

    def __init__(self):
        self.cloudpickle = importlib.import_module("cloudpickle")

    def serialize(self, writer: ModelStateWriter, obj: object):
        with writer.open(self.file_name, "wb") as cloudpickle_file:
            self.cloudpickle.dump(obj, cloudpickle_file)

    def deserialize(self, reader: ModelStateReader) -> object:
        with reader.open(self.file_name, "rb") as cloudpickle_file:
            obj = self.cloudpickle.load(cloudpickle_file)
        return obj

DillSerializer

dill.pkl을 사용합니다. 저장할 때 recurse=True로 객체가 참조하는 상태를 함께 처리하고 로드 결과를 반환합니다.

import importlib
from types import ModuleType
from palantir_models.models._serialization import ModelSerializer
from palantir_models.models._state_accessors import ModelStateWriter, ModelStateReader


class DillSerializer(ModelSerializer[object]):
    """Serializer utilizing the dill library for generic objects"""

    file_name = "dill.pkl"
    dill: ModuleType

    def __init__(self):
        self.dill = importlib.import_module("dill")

    def serialize(self, writer: ModelStateWriter, obj: object):
        with writer.open(self.file_name, "wb") as dill_file:
            self.dill.dump(obj, dill_file, recurse=True)

    def deserialize(self, reader: ModelStateReader) -> object:
        with reader.open(self.file_name, "rb") as dill_file:
            obj = self.dill.load(dill_file)
        return obj

HfAutoModelSerializer

HfAutoModelSerializer(model_class=None, **load_kwargs)는 기본으로 transformers.AutoModel을 사용합니다. AutoModelForSequenceClassification이나 BertForTokenClassification 같은 구체 클래스를 지정할 수 있습니다. model 디렉터리에 save_pretrained로 저장하고 from_pretrained에 저장 경로 및 load_kwargs를 전달합니다. num_labels=2 같은 로드 옵션도 이 경로로 전달합니다.

import importlib
from palantir_models import ModelSerializer, ModelStateReader, ModelStateWriter

class HfAutoModelSerializer(ModelSerializer):
    """
    Serializer for huggingface transformers AutoModel classes, using the
    from_pretrained and save_pretrained methods.
    Allows configuring a specific subclass (e.g. AutoModelForSequenceClassification or
    BertForTokenClassification) and passing additional kwargs to from_pretrained
    (e.g. num_labels=2).
    """

    DIR_NAME = "model"

    def __init__(self, model_class=None, **load_kwargs):
        if model_class is None:
            transformers = importlib.import_module("transformers")
            model_class = transformers.AutoModel
        self.model_class = model_class
        self.load_kwargs = load_kwargs

    def serialize(self, writer: ModelStateWriter, obj):
        model_dir = writer.mkdir(self.DIR_NAME)
        obj.save_pretrained(model_dir)

    def deserialize(self, reader: ModelStateReader):
        model_dir = reader.dir(self.DIR_NAME)
        return self.model_class.from_pretrained(model_dir, **self.load_kwargs)

HfAutoTokenizerSerializer

HfAutoTokenizerSerializer(tokenizer_class=None, **load_kwargs)는 기본 AutoTokenizer를 사용합니다. tokenizer 디렉터리에 저장하며 구체 토크나이저 클래스와 from_pretrained 옵션을 지정할 수 있습니다.

import importlib
from palantir_models import ModelSerializer, ModelStateReader, ModelStateWriter

class HfAutoTokenizerSerializer(ModelSerializer):
    """
    Serializer for huggingface transformers AutoTokenizer.
    """

    DIR_NAME = "tokenizer"

    def __init__(self, tokenizer_class=None, **load_kwargs):
        if tokenizer_class is None:
            transformers = importlib.import_module("transformers")
            tokenizer_class = transformers.AutoTokenizer
        self.tokenizer_class = tokenizer_class
        self.load_kwargs = load_kwargs

    def serialize(self, writer: ModelStateWriter, obj):
        tokenizer_dir = writer.mkdir(self.DIR_NAME)
        obj.save_pretrained(tokenizer_dir)

    def deserialize(self, reader: ModelStateReader):
        tokenizer_dir = reader.dir(self.DIR_NAME)
        return self.tokenizer_class.from_pretrained(tokenizer_dir, **self.load_kwargs)

HfPipelineSerializer

HfPipelineSerializer(pipeline_type, **load_kwargs)는 필수 작업 이름과 추가 로드 옵션을 받습니다. pipeline 디렉터리에 저장하고 transformers.pipeline에 작업 이름, model 경로와 키워드 인자를 전달하여 다시 생성합니다.

import importlib
from palantir_models import ModelSerializer, ModelStateReader, ModelStateWriter


class HfPipelineSerializer(ModelSerializer):
    """
    Serializer for huggingface transformers pipelines.
    Allows setting the pipeline task (e.g. sentiment-analysis).
    """

    DIR_NAME = "pipeline"

    def __init__(self, pipeline_type, **load_kwargs):
        self.transformers = importlib.import_module("transformers")
        self.pipeline_type = pipeline_type
        self.load_kwargs = load_kwargs

    def serialize(self, writer: ModelStateWriter, obj):
        pipeline_dir = writer.mkdir(self.DIR_NAME)
        obj.save_pretrained(pipeline_dir)

    def deserialize(self, reader: ModelStateReader):
        pipeline_dir = reader.dir(self.DIR_NAME)
        return self.transformers.pipeline(self.pipeline_type, model=pipeline_dir, **self.load_kwargs)

JsonSerializer

config.json을 텍스트 모드로 열어 json.dump 및 json.load로 처리합니다. JSON으로 표현 가능한 객체와 사전에 사용합니다.

import importlib
from types import ModuleType
from typing import Dict
from palantir_models.models._serialization import ModelSerializer
from palantir_models.models._state_accessors import ModelStateWriter, ModelStateReader


class JsonSerializer(ModelSerializer[Dict]):
    """Serializer for json-convertible objects and dictionaries"""

    file_name = "config.json"
    json: ModuleType

    def __init__(self):
        self.json = importlib.import_module("json")

    def serialize(self, writer: ModelStateWriter, obj: Dict):
        with writer.open(self.file_name, "w") as conf:
            self.json.dump(obj, conf)

    def deserialize(self, reader: ModelStateReader) -> Dict:
        with reader.open(self.file_name, "r") as conf:
            return self.json.load(conf)

PytorchStateSerializer

model_state_dict.pt에 PyTorch 상태 사전을 바이너리로 저장합니다. deserialize는 torch.load의 사전을 반환합니다. 신경망 클래스 인스턴스를 새로 만드는 동작은 포함하지 않습니다.

import importlib
from types import ModuleType
from palantir_models.models._serialization import ModelSerializer
from palantir_models.models._state_accessors import ModelStateWriter, ModelStateReader


class PytorchStateSerializer(ModelSerializer):
    """Serializer for PyTorch state dictionaries."""

    STATE_DICT_FILE_NAME = "model_state_dict.pt"
    torch: ModuleType

    def __init__(self):
        self.torch = importlib.import_module("torch")

    def serialize(self, writer: ModelStateWriter, obj: dict):
        """Serializes the state_dict of a PyTorch model."""
        with writer.open(self.STATE_DICT_FILE_NAME, "wb") as file_path:
            self.torch.save(obj, file_path)

    def deserialize(self, reader: ModelStateReader) -> dict:
        """Deserializes the state_dict of a PyTorch model."""
        with reader.open(self.STATE_DICT_FILE_NAME, "rb") as file_path:
            state_dict = self.torch.load(file_path)
            return state_dict

TensorflowKerasSerializer

TensorflowKerasSerializer(format=TensorflowFormat.DIR, custom_objects=None)는 tensorflow_saved_model_dir 아래에 모델을 저장합니다.

load_model에 custom_objects와 compile=False를 전달한 뒤 복원 객체의 compile()을 호출합니다. save에서 Invalid filepath extension for saving 오류가 발생하면 생성자의 형식을 DIR, H5 또는 KERAS 중 맞는 값으로 지정하라는 ValueError를 발생시킵니다. 다른 ValueError는 원래 오류를 그대로 전달합니다.

import enum
import importlib
import os
from types import ModuleType
from typing import Any, Dict, Optional

from palantir_models.models._serialization import ModelSerializer
from palantir_models.models._state_accessors import ModelStateReader, ModelStateWriter


class TensorflowFormat(enum.Enum):
    DIR = 0
    H5 = 1
    KERAS = 2

    def get_save_path(self, dir_path):
        if self == TensorflowFormat.DIR:
            return dir_path
        if self == TensorflowFormat.H5:
            return os.path.join(dir_path, "model.h5")
        if self == TensorflowFormat.KERAS:
            return os.path.join(dir_path, "model.keras")


class TensorflowKerasSerializer(ModelSerializer):
    """Serializer for tensorflow keras models"""

    DIR_NAME: str = "tensorflow_saved_model_dir"
    __tensorflow: ModuleType

    def __init__(self, format=TensorflowFormat.DIR, custom_objects: Optional[Dict[str, Any]] = None):
        self.__tensorflow = importlib.import_module("tensorflow")
        self.__custom_objects = custom_objects
        self.__format = format

    def serialize(self, writer: ModelStateWriter, obj: "tensorflow.keras.Model"):
        dir_path = writer.mkdir(self.DIR_NAME)
        self._save_model(dir_path, obj)

    def deserialize(self, reader: ModelStateReader) -> "tensorflow.keras.Model":
        dir_path = reader.dir(self.DIR_NAME)
        obj = self._load_model(dir_path)
        obj.compile()
        return obj

    def _save_model(self, dir_path: str, obj: "tensorflow.keras.Model"):
        try:
            save_path = self.__format.get_save_path(dir_path)
            obj.save(save_path)
        except ValueError as exc:
            if "Invalid filepath extension for saving" in str(exc):
                raise ValueError(
                    "Serialization failed due to invalid save format. Specify the correct file format in the TensorflowKerasSerializer constructor:\n"
                    "\tTensorflowKerasSerializer(format=TensorflowFormat.DIR) to use directory saving. This save format is considered legacy in Tensorflow\n"
                    "\tTensorflowKerasSerializer(format=TensorflowFormat.H5) to use .h5 file format saving\n"
                    "\tTensorflowKerasSerializer(format=TensorflowFormat.KERAS) to use .keras file format saving\n"
                ) from exc
            raise exc

    def _load_model(self, dir_path: str) -> "tensorflow.keras.Model":
        save_path = self.__format.get_save_path(dir_path)
        return self.__tensorflow.keras.models.load_model(save_path, custom_objects=self.__custom_objects, compile=False)

XGBoostSerializer

XGBModel을 xgboost_model.json에 저장합니다. 로드 시 XGBModel()을 만든 뒤 load_model로 파일을 읽어 반환합니다. 저장과 로드는 writer 및 reader가 제공한 파일의 실제 name 경로를 사용합니다.

from palantir_models import ModelSerializer
from palantir_models.models._state_accessors import ModelStateReader, ModelStateWriter
from xgboost.sklearn import XGBModel


class XGBoostSerializer(ModelSerializer[XGBModel]):
    """Simple Serializer for XGBoost SkLearn Models."""

    file_name = "xgboost_model.json"

    def serialize(self, writer: ModelStateWriter, obj: XGBModel):
        with writer.open(self.file_name, "w") as xgbfile:
            obj.save_model(xgbfile.name)

    def deserialize(self, reader: ModelStateReader) -> XGBModel:
        model = XGBModel()
        with reader.open(self.file_name, "r") as xgbfile:
            model.load_model(xgbfile.name)
            return model

YamlSerializer

config.yaml을 텍스트 모드로 열어 safe_dump로 저장하고 safe_load로 복원합니다. YAML로 표현 가능한 객체와 사전에 사용합니다.

import importlib
from types import ModuleType
from typing import Dict
from palantir_models.models._serialization import ModelSerializer
from palantir_models.models._state_accessors import ModelStateWriter, ModelStateReader


class YamlSerializer(ModelSerializer[Dict]):
    """Serializer for yaml-convertible objects and dictionaries"""

    file_name = "config.yaml"
    yaml: ModuleType

    def __init__(self):
        self.yaml = importlib.import_module("yaml")

    def serialize(self, writer: ModelStateWriter, obj: Dict):
        with writer.open(self.file_name, "w") as conf:
            self.yaml.safe_dump(obj, conf)

    def deserialize(self, reader: ModelStateReader) -> Dict:
        with reader.open(self.file_name, "r") as conf:
            return self.yaml.safe_load(conf)