AWS Trainium & Inferentia 文件
在 AWS Inferentia 上使用 Optimum Neuron 執行 Sentence Transformers
並獲得增強的文件體驗
開始使用
在 AWS Inferentia 上使用 Optimum Neuron 執行 Sentence Transformers
文字模型
本教程的筆記本版本可在此處找到：此處。
本指南解釋瞭如何在 AWS Inferentia2 上使用 Optimum Neuron 編譯、載入和使用 Sentence Transformers (SBERT) 模型，從而實現高效的嵌入計算。Sentence Transformers 是用於生成句子嵌入的強大模型。您可以使用此 Sentence Transformers 計算 100 多種語言的句子/文字嵌入。然後，這些嵌入可以與餘弦相似度等進行比較，以找到含義相似的句子。這對於語義文字相似度、語義搜尋或釋義挖掘非常有用。
將 Sentence Transformers 模型轉換為 AWS Inferentia2
首先，您需要將 Sentence Transformers 模型轉換為與 AWS Inferentia2 相容的格式。您可以使用 `optimum-cli` 或 `NeuronModelForSentenceTransformers` 類使用 Optimum Neuron 編譯 Sentence Transformers 模型。下面您將找到兩種方法的示例。我們必須確保已安裝 `sentence-transformers`。這僅在匯出模型時需要。
pip install sentence-transformers
在此處，我們將使用 `NeuronModelForSentenceTransformers`，它可以用於將任何 Sentence Transformers 模型轉換為與 AWS Inferentia2 相容的格式，或載入已轉換的模型。使用 `NeuronModelForSentenceTransformers` 匯出模型時，您需要設定 `export=True` 並定義輸入形狀和批處理大小。輸入形狀由 `sequence_length` 定義，批處理大小由 `batch_size` 定義。
from optimum.neuron import NeuronModelForSentenceTransformers
# Sentence Transformers model from HuggingFace
model_id = "BAAI/bge-small-en-v1.5"
input_shapes = {"batch_size": 1, "sequence_length": 384} # mandatory shapes
# Load Transformers model and export it to AWS Inferentia2
model = NeuronModelForSentenceTransformers.from_pretrained(model_id, export=True, **input_shapes)
# Save model to disk
model.save_pretrained("bge_emb_inf2/")在此處，我們將使用 `optimum-cli` 轉換模型。與 `NeuronModelForSentenceTransformers` 類似，我們需要定義輸入形狀和批處理大小。輸入形狀由 `sequence_length` 定義，批處理大小由 `batch_size` 定義。`optimum-cli` 將自動將模型轉換為與 AWS Inferentia2 相容的格式，並將其儲存到指定的輸出目錄。
optimum-cli export neuron -m BAAI/bge-small-en-v1.5 --sequence_length 384 --batch_size 1 --task feature-extraction bge_emb_inf2/載入已編譯的 Sentence Transformers 模型並執行推理
一旦我們擁有一個已編譯的 Sentence Transformers 模型（無論是我們自己匯出的還是 Hugging Face Hub 上可用的），我們就可以載入它並執行推理。為了載入模型，我們可以使用 `NeuronModelForSentenceTransformers` 類，它是 `SentenceTransformer` 類的一個抽象層。`NeuronModelForSentenceTransformers` 類將自動將輸入填充到指定的 `sequence_length`，並在 AWS Inferentia2 上執行推理。
from optimum.neuron import NeuronModelForSentenceTransformers
from transformers import AutoTokenizer
model_id_or_path = "bge_emb_inf2/"
tokenizer_id = "BAAI/bge-small-en-v1.5"
# Load model and tokenizer
model = NeuronModelForSentenceTransformers.from_pretrained(model_id_or_path)
tokenizer = AutoTokenizer.from_pretrained(tokenizer_id)
# Run inference
prompt = "I like to eat apples"
encoded_input = tokenizer(prompt, return_tensors='pt')
outputs = model(**encoded_input)
token_embeddings = outputs.token_embeddings
sentence_embedding = outputs.sentence_embedding
print(f"token embeddings: {token_embeddings.shape}") # torch.Size([1, 7, 384])
print(f"sentence_embedding: {sentence_embedding.shape}") # torch.Size([1, 384])生產用途
有關在生產環境中部署這些模型的資訊，請參閱 Amazon SageMaker 部落格。
CLIP
為 AWS Inferentia2 編譯 CLIP
您可以使用 `optimum-cli` 或 `NeuronModelForSentenceTransformers` 類使用 Optimum Neuron 編譯 CLIP 模型。選擇您喜歡的一種方法
- 使用 Optimum CLI
optimum-cli export neuron -m sentence-transformers/clip-ViT-B-32 --sequence_length 64 --text_batch_size 3 --image_batch_size 1 --num_channels 3 --height 224 --width 224 --task feature-extraction --subfolder 0_CLIPModel clip_emb/- 使用 `NeuronModelForSentenceTransformers` 類
from optimum.neuron import NeuronModelForSentenceTransformers
model_id = "sentence-transformers/clip-ViT-B-32"
# configs for compiling model
input_shapes = {
"num_channels": 3,
"height": 224,
"width": 224,
"text_batch_size": 3,
"image_batch_size": 1,
"sequence_length": 64,
}
emb_model = NeuronModelForSentenceTransformers.from_pretrained(
model_id, subfolder="0_CLIPModel", export=True, library_name="sentence_transformers", dynamic_batch_size=False, **input_shapes
)
# Save locally or upload to the HuggingFace Hub
save_directory = "clip_emb/"
emb_model.save_pretrained(save_directory)載入已編譯的 Sentence Transformers 模型並執行推理
from PIL import Image
from sentence_transformers import util
from transformers import CLIPProcessor
from optimum.neuron import NeuronModelForSentenceTransformers
save_directory = "clip_emb"
emb_model = NeuronModelForSentenceTransformers.from_pretrained(save_directory)
processor = CLIPProcessor.from_pretrained(save_directory)
inputs = processor(
text=["Two dogs in the snow", 'A cat on a table', 'A picture of London at night'], images=Image.open("two_dogs_in_snow.jpg"), return_tensors="pt", padding=True
)
outputs = emb_model(**inputs)
# Compute cosine similarities
cos_scores = util.cos_sim(outputs.image_embeds, outputs.text_embeds)
print(cos_scores)
# tensor([[0.3072, 0.1016, 0.1095]])注意事項
由於啟用動態批處理的編譯模型只接受具有相同批處理大小的輸入張量，因此如果輸入文字和影像具有不同的批處理大小，我們無法設定 `dynamic_batch_size=True`。由於 `NeuronModelForSentenceTransformers` 類會將輸入填充到編譯期間使用的批處理大小（`text_batch_size` 和 `image_batch_size`），因此您可以在編譯期間使用相對較大的批處理大小以獲得靈活性，但會犧牲一些計算資源。
例如，如果您想編碼 3、4 或 5 個文字和 1 張影像，您可以在編譯期間設定 `text_batch_size = 5 = max(3, 4, 5)` 和 `image_batch_size = 1`。