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`。