Batch Transform with Custom Infrastructure

Amazon SageMaker
1

Hi,

I’m attempting to deploy a model which is trained via PyTorch but saved in HuggingFace and uses a custom class which extends PreTrainedModel:

class BertClassifierConfig(BertConfig):
    model_type = "BertClassifier"
    problem_type = "single_label_classification"

class BertClassifier(PreTrainedModel):
    config_class = BertClassifierConfig

    def __init__(self, labels, dropout=0.1):
        ...

    def forward(self, input_id, mask):
        ...

I’m saving this model to HuggingFace and attempting to deploy it with Batch Transform in SageMaker, but am having problems. Specifically KeyError: 'BertClassifier' — which seems to be because the custom architecture isn’t defined.

Is there an easy way to address this problem? I’ve found a few solutions:

  • Creating a custom inference.py (I have no docker experience, so this could be pretty rough, not to mention I haven’t found any documentation/tutorials around creating an image which extends the HF image, just PyTorch)
  • Reconfiguring my architecture to extend nn.Module, skip HuggingFace, and use PyTorch-based batch transform (still requires a docker image? Unsure)
  • Reconfigure the architecture to not implement a custom architecture — I haven’t been able to get this to match the performance of the custom architecture, so this is really not preferred.

Appreciate any and all input — very new to SageMaker and only a few months experience with HuggingFace!