使用 SDK 在 Amazon SageMaker 上訓練和部署 Hugging Face 模型

本入門指南將向您展示如何使用 SDK 在 Amazon SageMaker 上快速使用 Hugging Face。瞭解如何在 SageMaker 上對預訓練的 🤗 Transformers 模型進行微調和部署，以完成二元文字分類任務。

安裝與設定

首先安裝必要的 Hugging Face 庫和 SageMaker。如果您尚未安裝 PyTorch，也需要安裝它。如果您在 SageMaker Studio 中執行此示例，它已安裝在筆記本核心中！

pip install "sagemaker>=2.140.0" "transformers==4.26.1" "datasets[s3]==2.10.1" --upgrade

如果您想在 SageMaker Studio 中執行此示例，請升級 ipywidgets 以使用 🤗 Datasets 庫並重新啟動核心

%%capture
import IPython
!conda install -c conda-forge ipywidgets -y
IPython.Application.instance().kernel.do_shutdown(True)

接下來，您應該設定環境：一個 SageMaker 會話和一個 S3 儲存桶。S3 儲存桶將儲存資料、模型和日誌。您將需要訪問具有所需許可權的 IAM 執行角色。

如果您打算在本地環境中使用 SageMaker，您需要自己提供 role。在此處瞭解更多設定方法。

⚠️ 執行角色僅在您在 SageMaker 中執行筆記本時可用。如果您嘗試在非 SageMaker 筆記本中執行 get_execution_role，您將收到區域錯誤。

import sagemaker

sess = sagemaker.Session()
sagemaker_session_bucket = None
if sagemaker_session_bucket is None and sess is not None:
    sagemaker_session_bucket = sess.default_bucket()

role = sagemaker.get_execution_role()
sess = sagemaker.Session(default_bucket=sagemaker_session_bucket)

預處理

🤗 Datasets 庫可以輕鬆下載和預處理用於訓練的資料集。下載並標記 IMDb 資料集

from datasets import load_dataset
from transformers import AutoTokenizer

# load dataset
train_dataset, test_dataset = load_dataset("imdb", split=["train", "test"])

# load tokenizer
tokenizer = AutoTokenizer.from_pretrained("distilbert/distilbert-base-uncased")

# create tokenization function
def tokenize(batch):
    return tokenizer(batch["text"], padding="max_length", truncation=True)

# tokenize train and test datasets
train_dataset = train_dataset.map(tokenize, batched=True)
test_dataset = test_dataset.map(tokenize, batched=True)

# set dataset format for PyTorch
train_dataset =  train_dataset.rename_column("label", "labels")
train_dataset.set_format("torch", columns=["input_ids", "attention_mask", "labels"])
test_dataset = test_dataset.rename_column("label", "labels")
test_dataset.set_format("torch", columns=["input_ids", "attention_mask", "labels"])

將資料集上傳到 S3 儲存桶

接下來，使用 🤗 Datasets S3 檔案系統實現將預處理後的資料集上傳到您的 S3 會話儲存桶

# save train_dataset to s3
training_input_path = f's3://{sess.default_bucket()}/{s3_prefix}/train'
train_dataset.save_to_disk(training_input_path)

# save test_dataset to s3
test_input_path = f's3://{sess.default_bucket()}/{s3_prefix}/test'
test_dataset.save_to_disk(test_input_path)

開始訓練任務

建立一個 Hugging Face Estimator 來處理端到端 SageMaker 訓練和部署。需要注意的最重要引數是：

entry_point 指的是微調指令碼，您可以在 train.py 檔案中找到它。
instance_type 指的是將啟動的 SageMaker 例項。完整的例項型別列表請參見此處。
hyperparameters 指的是模型將進行微調的訓練超引數。

from sagemaker.huggingface import HuggingFace

hyperparameters={
    "epochs": 1,                                       # number of training epochs
    "train_batch_size": 32,                            # training batch size
    "model_name":"distilbert/distilbert-base-uncased"  # name of pretrained model
}

huggingface_estimator = HuggingFace(
    entry_point="train.py",                 # fine-tuning script to use in training job
    source_dir="./scripts",                 # directory where fine-tuning script is stored
    instance_type="ml.p3.2xlarge",          # instance type
    instance_count=1,                       # number of instances
    role=role,                              # IAM role used in training job to access AWS resources (S3)
    transformers_version="4.36",             # Transformers version
    pytorch_version="2.1.0",                  # PyTorch version
    py_version="py310",                      # Python version
    hyperparameters=hyperparameters         # hyperparameters to use in training job
)

一行程式碼即可開始訓練

huggingface_estimator.fit({"train": training_input_path, "test": test_input_path})

部署模型

訓練任務完成後，透過呼叫 deploy() 並指定例項數量和例項型別來部署您的微調模型

predictor = huggingface_estimator.deploy(initial_instance_count=1,"ml.g4dn.xlarge")

對您的資料呼叫 predict()

sentiment_input = {"inputs": "It feels like a curtain closing...there was an elegance in the way they moved toward conclusion. No fan is going to watch and feel short-changed."}

predictor.predict(sentiment_input)

執行請求後，刪除端點

predictor.delete_endpoint()

下一步是什麼？

恭喜您，您剛剛在 SageMaker 上對預訓練的 🤗 Transformers 模型進行了微調和部署，以完成二元文字分類！🎉

< > 在 GitHub 上更新