AWS Trainium 和 Inferentia 文档
在 AWS Inferentia 上使用 Optimum Neuron 的 Sentence Transformers
并获得增强的文档体验
开始使用
在 AWS Inferentia 上使用 Optimum Neuron 的 Sentence Transformers
文本模型
该教程有一个 notebook 版本，请点击这里。
本指南解释了如何在 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`。