AWS Trainium 和 Inferentia 文档

模型

AWS Trainium 和 Inferentia

加入 Hugging Face 社区

并获得增强的文档体验

在模型、数据集和空间上协作

通过加速推理获得更快的示例

在文档主题之间切换

开始使用

模型

通用模型类

NeuronTracedModel

NeuronTracedModel 类可用于实例化没有特定头的基本 Neuron 模型。它用作除文本生成之外的所有任务的基类。

class optimum.neuron.NeuronTracedModel

< 源代码 >

( model: ScriptModule config: PretrainedConfig model_save_dir: typing.Union[str, pathlib.Path, tempfile.TemporaryDirectory, NoneType] = None model_file_name: typing.Optional[str] = None preprocessors: typing.Optional[typing.List] = None neuron_config: typing.Optional[ForwardRef('NeuronDefaultConfig')] = None **kwargs )

在 Neuron 设备上运行编译和优化模型的基类。

它实现了与 Hugging Face Hub 交互的通用方法，以及将 vanilla transformers 模型编译为 neuron 优化的 TorchScript 模块，并使用 optimum.exporters.neuron 工具链导出它。

类属性

model_type (str, optional, defaults to "neuron_model") — 注册 NeuronTracedModel 类时要使用的模型类型名称。
auto_model_class (Type, optional, defaults to AutoModel) — 要由当前 NeuronTracedModel 类表示的 AutoModel 类。

常用属性

model (torch.jit._script.ScriptModule) — 为 neuron 设备编译的已加载 ScriptModule。
config (PretrainedConfig) — 模型的配置。
model_save_dir (Path) — 保存 neuron 编译模型的目录。默认情况下，如果加载的模型是本地的，则将使用原始模型所在的目录。否则，将使用缓存目录。

can_generate

< 源代码 >

( )

返回此模型是否可以使用 .generate() 生成序列。

get_input_static_shapes

< 源代码 >

( neuron_config: NeuronDefaultConfig )

获取包含有效静态形状的输入字典。

load_model

< 源代码 >

( path: typing.Union[str, pathlib.Path] to_neuron: bool = False device_id: int = 0 )

参数

path (Union[str, Path]) — 编译模型的路径。
to_neuron (bool, 默认为 False) — 是否手动将跟踪模型移动到 NeuronCore。仅当 inline_weights_to_neff=False 时才需要，否则它会自动加载到 Neuron 设备。
device_id (int, 默认为 0) — 将跟踪模型加载到的 NeuronCore 的索引。

加载由 neuron(x)-cc 编译器编译的 TorchScript 模块。它将首先加载到 CPU 上，然后移动到一个或多个 NeuronCore。

remove_padding

< 源代码 >

( outputs: typing.List[torch.Tensor] dims: typing.List[int] indices: typing.List[int] padding_side: typing.Literal['right', 'left'] = 'right' )

参数

outputs (List[torch.Tensor]) — 作为推理输出的 torch 张量列表。
dims (List[int]) — 我们在其中切片张量的维度列表。
indices (List[int]) — 我们在其中沿轴切片张量的索引列表。
padding_side (Literal["right", "left"], 默认为 “right”) — 已应用填充的一侧。

从输出张量中删除填充。

NeuronDecoderModel

NeuronDecoderModel 类是文本生成模型的基础类。

class optimum.neuron.NeuronDecoderModel

< 源代码 >

( config: PretrainedConfig checkpoint_dir: typing.Union[str, pathlib.Path, tempfile.TemporaryDirectory] compiled_dir: typing.Union[str, pathlib.Path, tempfile.TemporaryDirectory, NoneType] = None generation_config: typing.Optional[transformers.generation.configuration_utils.GenerationConfig] = None )

在 Neuron 设备上转换和运行预训练 transformers 解码器模型的基础类。

它实现了将预训练 transformers 解码器模型转换为 Neuron transformer 模型的方法，通过

将原始模型的检查点权重传输到优化的 neuron 图中，
使用 Neuron 编译器编译结果图。

常用属性

model (torch.nn.Module) — 具有为 neuron 设备优化的图的解码器模型。
config (PretrainedConfig) — 原始模型的配置。
generation_config (GenerationConfig) — 调用 generate() 时默认使用的生成配置。

自然语言处理

以下 Neuron 模型类可用于自然语言处理任务。

NeuronModelForFeatureExtraction

class optimum.neuron.NeuronModelForFeatureExtraction

< 源代码 >

参数

config (transformers.PretrainedConfig) — PretrainedConfig 是模型配置类，包含模型的所有参数。使用配置文件初始化不会加载与模型关联的权重，只会加载配置。查看 optimum.neuron.modeling.NeuronTracedModel.from_pretrained 方法以加载模型权重。
model (torch.jit._script.ScriptModule) — torch.jit._script.ScriptModule 是 TorchScript 模块，其中嵌入了由 neuron(x) 编译器编译的 NEFF（Neuron 可执行文件格式）。

具有用于特征提取任务的 BaseModelOutput 的 Neuron 模型。

此模型继承自 ~neuron.modeling.NeuronTracedModel。查看超类文档，了解库为所有模型实现的通用方法（例如下载或保存）

Neuron 设备上的特征提取模型。

forward

< source >

( input_ids: Tensor attention_mask: Tensor token_type_ids: typing.Optional[torch.Tensor] = None **kwargs )

参数

input_ids (形状为 (batch_size, sequence_length) 的 torch.Tensor) — 词汇表中输入序列 tokens 的索引。索引可以使用 AutoTokenizer 获得。有关详细信息，请参阅 PreTrainedTokenizer.encode 和 PreTrainedTokenizer.__call__。什么是输入 IDs?
attention_mask (形状为 `(batch_size, sequence_length)` 的 `Union[torch.Tensor, None]`, 默认为 `None`) — 用于避免对 padding token 索引执行 attention 的 Mask。 Mask 值在 [0, 1] 中选择：
- 1 代表 **未被 mask** 的 tokens，
- 0 代表 **已被 mask** 的 tokens。什么是 attention masks?
token_type_ids (形状为 `(batch_size, sequence_length)` 的 `Union[torch.Tensor, None]`, 默认为 `None`) — 用于指示输入的第一部分和第二部分的 Segment token 索引。索引在 [0, 1] 中选择：
- 1 代表 **句子 A** 的 tokens，
- 0 代表 **句子 B** 的 tokens。什么是 token type IDs?

NeuronModelForFeatureExtraction 的 forward 方法，会覆盖 __call__ 特殊方法。

尽管 forward pass 的配方需要在该函数中定义，但应该在之后调用 Module 实例而不是此函数，因为前者会负责运行预处理和后处理步骤，而后者会默默地忽略它们。

特征提取示例： _（以下模型使用 neuronx 编译器编译，只能在 INF2 上运行。如果您使用的是 INF1，请将 “neuronx” 替换为 “neuron”。）_

>>> from transformers import AutoTokenizer
>>> from optimum.neuron import NeuronModelForFeatureExtraction

>>> tokenizer = AutoTokenizer.from_pretrained("optimum/all-MiniLM-L6-v2-neuronx")
>>> model = NeuronModelForFeatureExtraction.from_pretrained("optimum/all-MiniLM-L6-v2-neuronx")

>>> inputs = tokenizer("Dear Evan Hansen is the winner of six Tony Awards.", return_tensors="pt")

>>> outputs = model(**inputs)
>>> last_hidden_state = outputs.last_hidden_state
>>> list(last_hidden_state.shape)
[1, 13, 384]

NeuronModelForSentenceTransformers

class optimum.neuron.NeuronModelForSentenceTransformers

< source >

参数

config (transformers.PretrainedConfig) — PretrainedConfig 是模型配置类，其中包含模型的所有参数。使用配置文件初始化不会加载与模型关联的权重，只会加载配置。查看 optimum.neuron.modeling.NeuronTracedModel.from_pretrained 方法以加载模型权重。
model (torch.jit._script.ScriptModule) — torch.jit._script.ScriptModule 是 TorchScript 模块，其中嵌入了由 neuron(x) 编译器编译的 NEFF（Neuron Executable File Format）文件。

用于 Sentence Transformers 的 Neuron 模型。

此模型继承自 ~neuron.modeling.NeuronTracedModel。查看超类文档，了解库为所有模型实现的通用方法（例如下载或保存）

Neuron 设备上的 Sentence Transformers 模型。

forward

< source >

( input_ids: Tensor attention_mask: Tensor pixel_values: typing.Optional[torch.Tensor] = None token_type_ids: typing.Optional[torch.Tensor] = None **kwargs )

参数

input_ids (形状为 `(batch_size, sequence_length)` 的 `torch.Tensor`) — 词汇表中输入序列 tokens 的索引。索引可以使用 AutoTokenizer 获得。有关详细信息，请参阅 PreTrainedTokenizer.encode 和 PreTrainedTokenizer.__call__。什么是输入 IDs?
attention_mask (形状为 `(batch_size, sequence_length)` 的 `Union[torch.Tensor, None]`, 默认为 `None`) — 用于避免对 padding token 索引执行 attention 的 Mask。 Mask 值在 [0, 1] 中选择：
- 1 代表 **未被 mask** 的 tokens，
- 0 代表 **已被 mask** 的 tokens。什么是 attention masks?
token_type_ids (形状为 `(batch_size, sequence_length)` 的 `Union[torch.Tensor, None]`, 默认为 `None`) — 用于指示输入的第一部分和第二部分的 Segment token 索引。索引在 [0, 1] 中选择：
- 1 代表 **句子 A** 的 tokens，
- 0 代表 **句子 B** 的 tokens。什么是 token type IDs?

NeuronModelForSentenceTransformers 的 forward 方法，会覆盖 __call__ 特殊方法。

TEXT Sentence Transformers 示例

>>> from transformers import AutoTokenizer
>>> from optimum.neuron import NeuronModelForSentenceTransformers

>>> tokenizer = AutoTokenizer.from_pretrained("optimum/bge-base-en-v1.5-neuronx")
>>> model = NeuronModelForSentenceTransformers.from_pretrained("optimum/bge-base-en-v1.5-neuronx")

>>> inputs = tokenizer("In the smouldering promise of the fall of Troy, a mythical world of gods and mortals rises from the ashes.", return_tensors="pt")

>>> outputs = model(**inputs)
>>> token_embeddings = outputs.token_embeddings
>>> sentence_embedding = = outputs.sentence_embedding

NeuronModelForMaskedLM

class optimum.neuron.NeuronModelForMaskedLM

< source >

参数

config (transformers.PretrainedConfig) — PretrainedConfig 是模型配置类，其中包含模型的所有参数。使用配置文件初始化不会加载与模型关联的权重，只会加载配置。查看 optimum.neuron.modeling.NeuronTracedModel.from_pretrained 方法以加载模型权重。
model (torch.jit._script.ScriptModule) — torch.jit._script.ScriptModule 是 TorchScript 模块，其中嵌入了由 neuron(x) 编译器编译的 NEFF（Neuron Executable File Format）文件。

具有 MaskedLMOutput 的 Neuron 模型，用于 masked language modeling 任务。

此模型继承自 ~neuron.modeling.NeuronTracedModel。查看超类文档，了解库为所有模型实现的通用方法（例如下载或保存）

用于 Neuron 设备的 masked language model。

forward

< source >

( input_ids: Tensor attention_mask: Tensor token_type_ids: typing.Optional[torch.Tensor] = None **kwargs )

参数

input_ids (形状为 `(batch_size, sequence_length)` 的 `torch.Tensor`) — 词汇表中输入序列 tokens 的索引。索引可以使用 AutoTokenizer 获得。有关详细信息，请参阅 PreTrainedTokenizer.encode 和 PreTrainedTokenizer.__call__。什么是输入 IDs?
attention_mask (形状为 `(batch_size, sequence_length)` 的 `Union[torch.Tensor, None]`, 默认为 `None`) — 用于避免对 padding token 索引执行 attention 的 Mask。 Mask 值在 [0, 1] 中选择：
- 1 代表 **未被 mask** 的 tokens，
- 0 代表 **已被 mask** 的 tokens。什么是 attention masks?
token_type_ids (形状为 `(batch_size, sequence_length)` 的 `Union[torch.Tensor, None]`, 默认为 `None`) — 用于指示输入的第一部分和第二部分的 Segment token 索引。索引在 [0, 1] 中选择：
- 1 代表 **句子 A** 的 tokens，
- 0 代表 **句子 B** 的 tokens。什么是 token type IDs?

NeuronModelForMaskedLM 的 forward 方法，会覆盖 __call__ 特殊方法。

fill mask 示例： _（以下模型使用 neuronx 编译器编译，只能在 INF2 上运行。如果您使用的是 INF1，请将 “neuronx” 替换为 “neuron”。）_

>>> from transformers import AutoTokenizer
>>> from optimum.neuron import NeuronModelForMaskedLM
>>> import torch

>>> tokenizer = AutoTokenizer.from_pretrained("optimum/legal-bert-base-uncased-neuronx")
>>> model = NeuronModelForMaskedLM.from_pretrained("optimum/legal-bert-base-uncased-neuronx")

>>> inputs = tokenizer("This [MASK] Agreement is between General Motors and John Murray.", return_tensors="pt")

>>> outputs = model(**inputs)
>>> logits = outputs.logits
>>> list(logits.shape)
[1, 13, 30522]

NeuronModelForSequenceClassification

class optimum.neuron.NeuronModelForSequenceClassification

< source >

参数

config (transformers.PretrainedConfig) — PretrainedConfig 是模型配置类，其中包含模型的所有参数。使用配置文件初始化不会加载与模型关联的权重，只会加载配置。查看 optimum.neuron.modeling.NeuronTracedModel.from_pretrained 方法以加载模型权重。
model (torch.jit._script.ScriptModule) — torch.jit._script.ScriptModule 是 TorchScript 模块，其中嵌入了由 neuron(x) 编译器编译的 NEFF（Neuron Executable File Format）文件。

具有序列分类/回归头的 Neuron 模型（池化输出顶部的线性层），例如用于 GLUE 任务。

此模型继承自 ~neuron.modeling.NeuronTracedModel。查看超类文档，了解库为所有模型实现的通用方法（例如下载或保存）

Neuron 设备上的序列分类模型。

forward

< source >

( input_ids: Tensor attention_mask: Tensor token_type_ids: typing.Optional[torch.Tensor] = None **kwargs )

参数

input_ids (torch.Tensor，形状为 (batch_size, sequence_length)) — 词汇表中输入序列 tokens 的索引。索引可以使用 AutoTokenizer 获取。有关详细信息，请参阅 PreTrainedTokenizer.encode 和 PreTrainedTokenizer.__call__。什么是输入 ID？
attention_mask (Union[torch.Tensor, None]，形状为 (batch_size, sequence_length)，默认为 None) — 用于避免在 padding token 索引上执行 attention 的掩码。掩码值在 [0, 1] 中选择：
- 1 表示 未被掩盖 的 tokens，
- 0 表示 被掩盖 的 tokens。什么是 attention masks？
token_type_ids (Union[torch.Tensor, None]，形状为 (batch_size, sequence_length)，默认为 None) — 分段 token 索引，用于指示输入的第一部分和第二部分。索引在 [0, 1] 中选择：
- 1 表示 句子 A 的 tokens，
- 0 表示 句子 B 的 tokens。什么是 token type IDs？

NeuronModelForSequenceClassification forward 方法，覆盖了 __call__ 特殊方法。

单标签分类示例： （以下模型使用 neuronx 编译器编译，并且只能在 INF2 上运行。）

>>> from transformers import AutoTokenizer
>>> from optimum.neuron import NeuronModelForSequenceClassification

>>> tokenizer = AutoTokenizer.from_pretrained("optimum/distilbert-base-uncased-finetuned-sst-2-english-neuronx")
>>> model = NeuronModelForSequenceClassification.from_pretrained("optimum/distilbert-base-uncased-finetuned-sst-2-english-neuronx")

>>> inputs = tokenizer("Hamilton is considered to be the best musical of human history.", return_tensors="pt")

>>> outputs = model(**inputs)
>>> logits = outputs.logits
>>> list(logits.shape)
[1, 2]

NeuronModelForQuestionAnswering

class optimum.neuron.NeuronModelForQuestionAnswering

< source >

参数

config (transformers.PretrainedConfig) — PretrainedConfig 是模型配置类，包含模型的所有参数。使用配置文件初始化不会加载与模型关联的权重，只会加载配置。查看 optimum.neuron.modeling.NeuronTracedModel.from_pretrained 方法以加载模型权重。
model (torch.jit._script.ScriptModule) — torch.jit._script.ScriptModule 是 TorchScript 模块，其中嵌入了由 neuron(x) 编译器编译的 NEFF（神经元可执行文件格式）。

具有 QuestionAnsweringModelOutput 的 Neuron 模型，用于抽取式问答任务，如 SQuAD。

此模型继承自 ~neuron.modeling.NeuronTracedModel。查看超类文档，了解库为所有模型实现的通用方法（例如下载或保存）

Neuron 设备上的问答模型。

forward

< source >

( input_ids: Tensor attention_mask: Tensor token_type_ids: typing.Optional[torch.Tensor] = None **kwargs )

参数

input_ids (torch.Tensor，形状为 (batch_size, sequence_length)) — 词汇表中输入序列 tokens 的索引。索引可以使用 AutoTokenizer 获取。有关详细信息，请参阅 PreTrainedTokenizer.encode 和 PreTrainedTokenizer.__call__。什么是输入 ID？
attention_mask (Union[torch.Tensor, None]，形状为 (batch_size, sequence_length)，默认为 None) — 用于避免在 padding token 索引上执行 attention 的掩码。掩码值在 [0, 1] 中选择：
- 1 表示 未被掩盖 的 tokens，
- 0 表示 被掩盖 的 tokens。什么是 attention masks？
token_type_ids (Union[torch.Tensor, None]，形状为 (batch_size, sequence_length)，默认为 None) — 分段 token 索引，用于指示输入的第一部分和第二部分。索引在 [0, 1] 中选择：
- 1 表示 句子 A 的 tokens，
- 0 表示 句子 B 的 tokens。什么是 token type IDs？

NeuronModelForQuestionAnswering forward 方法，覆盖了 __call__ 特殊方法。

问答示例： （以下模型使用 neuronx 编译器编译，并且只能在 INF2 上运行。）

>>> import torch
>>> from transformers import AutoTokenizer
>>> from optimum.neuron import NeuronModelForQuestionAnswering

>>> tokenizer = AutoTokenizer.from_pretrained("optimum/roberta-base-squad2-neuronx")
>>> model = NeuronModelForQuestionAnswering.from_pretrained("optimum/roberta-base-squad2-neuronx")

>>> question, text = "Are there wheelchair spaces in the theatres?", "Yes, we have reserved wheelchair spaces with a good view."
>>> inputs = tokenizer(question, text, return_tensors="pt")
>>> start_positions = torch.tensor([1])
>>> end_positions = torch.tensor([12])

>>> outputs = model(**inputs, start_positions=start_positions, end_positions=end_positions)
>>> start_scores = outputs.start_logits
>>> end_scores = outputs.end_logits

NeuronModelForTokenClassification

class optimum.neuron.NeuronModelForTokenClassification

< source >

参数

config (transformers.PretrainedConfig) — PretrainedConfig 是模型配置类，包含模型的所有参数。使用配置文件初始化不会加载与模型关联的权重，只会加载配置。查看 optimum.neuron.modeling.NeuronTracedModel.from_pretrained 方法以加载模型权重。
model (torch.jit._script.ScriptModule) — torch.jit._script.ScriptModule 是 TorchScript 模块，其中嵌入了由 neuron(x) 编译器编译的 NEFF（神经元可执行文件格式）。

具有 token 分类头的 Neuron 模型（隐藏状态输出之上的线性层），例如用于命名实体识别 (NER) 任务。

此模型继承自 ~neuron.modeling.NeuronTracedModel。查看超类文档，了解库为所有模型实现的通用方法（例如下载或保存）

Neuron 设备上的 Token 分类模型。

forward

< source >

( input_ids: Tensor attention_mask: Tensor token_type_ids: typing.Optional[torch.Tensor] = None **kwargs )

参数

input_ids (torch.Tensor，形状为 (batch_size, sequence_length)) — 词汇表中输入序列 tokens 的索引。索引可以使用 AutoTokenizer 获取。有关详细信息，请参阅 PreTrainedTokenizer.encode 和 PreTrainedTokenizer.__call__。什么是输入 ID？
attention_mask (Union[torch.Tensor, None]，形状为 (batch_size, sequence_length)，默认为 None) — 用于避免在 padding token 索引上执行 attention 的掩码。掩码值在 [0, 1] 中选择：
- 1 表示 未被掩盖 的 tokens，
- 0 表示 被掩盖 的 tokens。什么是 attention masks？
token_type_ids (Union[torch.Tensor, None]，形状为 (batch_size, sequence_length)，默认为 None) — 分段 token 索引，用于指示输入的第一部分和第二部分。索引在 [0, 1] 中选择：
- 1 表示 句子 A 的 tokens，
- 0 表示 句子 B 的 tokens。什么是 token type IDs？

NeuronModelForTokenClassification forward 方法，覆盖了 __call__ 特殊方法。

token 分类示例： （以下模型使用 neuronx 编译器编译，并且只能在 INF2 上运行。）

>>> from transformers import AutoTokenizer
>>> from optimum.neuron import NeuronModelForTokenClassification

>>> tokenizer = AutoTokenizer.from_pretrained("optimum/bert-base-NER-neuronx")
>>> model = NeuronModelForTokenClassification.from_pretrained("optimum/bert-base-NER-neuronx")

>>> inputs = tokenizer("Lin-Manuel Miranda is an American songwriter, actor, singer, filmmaker, and playwright.", return_tensors="pt")

>>> outputs = model(**inputs)
>>> logits = outputs.logits
>>> list(logits.shape)
[1, 20, 9]

NeuronModelForMultipleChoice

class optimum.neuron.NeuronModelForMultipleChoice

< source >

参数

config (transformers.PretrainedConfig) — PretrainedConfig 是模型配置类，包含模型的所有参数。使用配置文件初始化不会加载与模型关联的权重，只会加载配置。查看 optimum.neuron.modeling.NeuronTracedModel.from_pretrained 方法以加载模型权重。
model (torch.jit._script.ScriptModule) — torch.jit._script.ScriptModule 是 TorchScript 模块，其中嵌入了由 neuron(x) 编译器编译的 NEFF（神经元可执行文件格式）。

具有多项选择分类头的 Neuron 模型（pooled 输出之上的线性层和一个 softmax），例如用于 RocStories/SWAG 任务。

此模型继承自 ~neuron.modeling.NeuronTracedModel。查看超类文档，了解库为所有模型实现的通用方法（例如下载或保存）

Neuron 设备上的多项选择模型。

forward

< source >

( input_ids: Tensor attention_mask: Tensor token_type_ids: typing.Optional[torch.Tensor] = None **kwargs )

参数

input_ids (torch.Tensor，形状为 (batch_size, num_choices, sequence_length)) — 词汇表中输入序列 tokens 的索引。索引可以使用 AutoTokenizer 获取。有关详细信息，请参阅 PreTrainedTokenizer.encode 和 PreTrainedTokenizer.__call__。什么是输入 ID？
attention_mask (Union[torch.Tensor, None]，形状为 (batch_size, num_choices, sequence_length)，默认为 None) — 用于避免在 padding token 索引上执行 attention 的掩码。掩码值在 [0, 1] 中选择：
- 1 表示 未被掩盖 的 tokens，
- 0 表示 被掩盖 的 tokens。什么是 attention masks？
token_type_ids (Union[torch.Tensor, None]，形状为 (batch_size, num_choices, sequence_length)，默认为 None) — 分段 token 索引，用于指示输入的第一部分和第二部分。索引在 [0, 1] 中选择：
- 1 表示 句子 A 的 tokens，
- 0 表示 句子 B 的 tokens。什么是 token type IDs？

NeuronModelForMultipleChoice forward 方法，覆盖了 __call__ 特殊方法。

多项选择示例： （以下模型使用 neuronx 编译器编译，并且只能在 INF2 上运行。）

>>> from transformers import AutoTokenizer
>>> from optimum.neuron import NeuronModelForMultipleChoice

>>> tokenizer = AutoTokenizer.from_pretrained("optimum/bert-base-uncased_SWAG-neuronx")
>>> model = NeuronModelForMultipleChoice.from_pretrained("optimum/bert-base-uncased_SWAG-neuronx")

>>> num_choices = 4
>>> first_sentence = ["Members of the procession walk down the street holding small horn brass instruments."] * num_choices
>>> second_sentence = [
...     "A drum line passes by walking down the street playing their instruments.",
...     "A drum line has heard approaching them.",
...     "A drum line arrives and they're outside dancing and asleep.",
...     "A drum line turns the lead singer watches the performance."
... ]
>>> inputs = tokenizer(first_sentence, second_sentence, truncation=True, padding=True)

# Unflatten the inputs values expanding it to the shape [batch_size, num_choices, seq_length]
>>> for k, v in inputs.items():
...     inputs[k] = [v[i: i + num_choices] for i in range(0, len(v), num_choices)]
>>> inputs = dict(inputs.convert_to_tensors(tensor_type="pt"))
>>> outputs = model(**inputs)
>>> logits = outputs.logits
>>> logits.shape
[1, 4]

NeuronModelForCausalLM

class optimum.neuron.NeuronModelForCausalLM

< source >

( config: PretrainedConfig checkpoint_dir: typing.Union[str, ForwardRef('Path'), ForwardRef('TemporaryDirectory')] compiled_dir: typing.Union[str, ForwardRef('Path'), ForwardRef('TemporaryDirectory'), NoneType] = None generation_config: typing.Optional[ForwardRef('GenerationConfig')] = None )

参数

model (torch.nn.Module) — torch.nn.Module 是 neuron 解码器图。
config (transformers.PretrainedConfig) — PretrainedConfig 是模型配置类，包含模型的所有参数。
model_path (Path) — 存储模型编译产物的目录。如果模型之前从未在本地保存过，则可以是临时目录。
generation_config (transformers.GenerationConfig) — GenerationConfig 包含模型生成任务的配置。

用于 Neuron 设备推理的带有因果语言建模头的 Neuron 模型。

此模型继承自 ~neuron.modeling.NeuronDecoderModel。查看超类文档，了解库为所有模型实现的通用方法（例如下载或保存）。

forward

< source >

( input_ids: Tensor cache_ids: Tensor start_ids: Tensor = None return_dict: bool = True )

参数

input_ids (torch.LongTensor) — 解码器输入序列标记在词汇表中的索引，形状为 (batch_size, sequence_length)。
cache_ids (torch.LongTensor) — 缓存的键和值需要在当前输入的索引处存储。
start_ids (torch.LongTensor) — 要处理的第一个标记的索引，从注意力掩码中推导出来。

NeuronModelForCausalLM 的 forward 方法，覆盖了 __call__ 特殊方法。

文本生成示例

>>> from transformers import AutoTokenizer
>>> from optimum.neuron import NeuronModelForCausalLM
>>> import torch

>>> tokenizer = AutoTokenizer.from_pretrained("gpt2")
>>> model = NeuronModelForCausalLM.from_pretrained("gpt2", export=True)

>>> inputs = tokenizer("My favorite moment of the day is", return_tensors="pt")

>>> gen_tokens = model.generate(**inputs, do_sample=True, temperature=0.9, min_length=20, max_length=20)
>>> tokenizer.batch_decode(gen_tokens)

NeuronModelForSeq2SeqLM

class optimum.neuron.NeuronModelForSeq2SeqLM

< source >

( encoder: ScriptModule decoder: ScriptModule config: PretrainedConfig model_save_dir: typing.Union[str, pathlib.Path, tempfile.TemporaryDirectory, NoneType] = None encoder_file_name: typing.Optional[str] = 'model.neuron' decoder_file_name: typing.Optional[str] = 'model.neuron' preprocessors: typing.Optional[typing.List] = None neuron_configs: typing.Optional[typing.Dict[str, ForwardRef('NeuronDefaultConfig')]] = None configs: typing.Optional[typing.Dict[str, ForwardRef('PretrainedConfig')]] = None generation_config: typing.Optional[transformers.generation.configuration_utils.GenerationConfig] = None **kwargs )

参数

encoder (torch.jit._script.ScriptModule) — torch.jit._script.ScriptModule 是编码器的 TorchScript 模块，其中嵌入了由 neuron(x) 编译器编译的 NEFF(Neuron 可执行文件格式)。
decoder (torch.jit._script.ScriptModule) — torch.jit._script.ScriptModule 是解码器的 TorchScript 模块，其中嵌入了由 neuron(x) 编译器编译的 NEFF(Neuron 可执行文件格式)。
config (transformers.PretrainedConfig) — PretrainedConfig 是模型配置类，包含模型的所有参数。使用配置文件初始化不会加载与模型关联的权重，仅加载配置。查看 optimum.neuron.modeling.NeuronTracedModel.from_pretrained 方法以加载模型权重。

用于 text2text-generation 任务的带有语言建模头的 Neuron 序列到序列模型。

此模型继承自 ~neuron.modeling.NeuronTracedModel。查看超类文档，了解库为所有模型实现的通用方法（例如下载或保存）

forward

< source >

( attention_mask: typing.Optional[torch.FloatTensor] = None decoder_input_ids: typing.Optional[torch.LongTensor] = None decoder_attention_mask: typing.Optional[torch.BoolTensor] = None encoder_outputs: typing.Optional[typing.Tuple[typing.Tuple[torch.Tensor]]] = None beam_scores: typing.Optional[torch.FloatTensor] = None return_dict: bool = False output_attentions: bool = False output_hidden_states: bool = False )

参数

input_ids (形状为 (batch_size, sequence_length) 的 torch.Tensor) — 词汇表中输入序列标记的索引。索引可以使用 AutoTokenizer 获得。有关详细信息，请参阅 PreTrainedTokenizer.encode 和 PreTrainedTokenizer.__call__。什么是输入 ID？
attention_mask (形状为 (batch_size, sequence_length) 的 Union[torch.Tensor, None], 默认为 None) — 用于避免在填充标记索引上执行注意力的掩码。在 [0, 1] 中选择的掩码值：
- 1 表示未被掩盖的标记，
- 0 表示被掩盖的标记。什么是注意力掩码？

NeuronModelForSeq2SeqLM 的 forward 方法，覆盖了 __call__ 特殊方法。

（以下模型使用 neuronx 编译器编译，只能在 INF2 上运行。）

使用小型 T5 模型进行文本到文本生成的示例

from transformers import AutoTokenizer
from optimum.neuron import NeuronModelForSeq2SeqLM

neuron_model = NeuronModelForSeq2SeqLM.from_pretrained(google-t5/t5-small, export=True, dynamic_batch_size=False, batch_size=1, sequence_length=64, num_beams=4)
neuron_model.save_pretrained("t5_small_neuronx")
del neuron_model

neuron_model = NeuronModelForSeq2SeqLM.from_pretrained("t5_small_neuronx")
tokenizer = AutoTokenizer.from_pretrained("t5_small_neuronx")
inputs = tokenizer("translate English to German: Lets eat good food.", return_tensors="pt")

output = neuron_model.generate(
    **inputs,
    num_return_sequences=1,
)
results = [tokenizer.decode(t, skip_special_tokens=True) for t in output]

（对于大型模型，为了适应 Neuron 核心，我们需要应用张量并行。以下是在 inf2.24xlarge 上运行的示例。）

使用张量并行进行文本到文本生成的示例

from transformers import AutoTokenizer
from optimum.neuron import NeuronModelForSeq2SeqLM
# 1. compile
if __name__ == "__main__":  # compulsory for parallel tracing since the API will spawn multiple processes.
    neuron_model = NeuronModelForSeq2SeqLM.from_pretrained(
        google/flan-t5-xl, export=True, tensor_parallel_size=8, dynamic_batch_size=False, batch_size=1, sequence_length=128, num_beams=4,
    )
    neuron_model.save_pretrained("flan_t5_xl_neuronx_tp8/")
    del neuron_model

# 2. inference
neuron_model = NeuronModelForSeq2SeqLM.from_pretrained("flan_t5_xl_neuronx_tp8")
tokenizer = AutoTokenizer.from_pretrained("flan_t5_xl_neuronx_tp8")
inputs = tokenizer("translate English to German: Lets eat good food.", return_tensors="pt")

output = neuron_model.generate(
    **inputs,
    num_return_sequences=1,
)
results = [tokenizer.decode(t, skip_special_tokens=True) for t in output]

计算机视觉

以下 Neuron 模型类可用于计算机视觉任务。

NeuronModelForImageClassification

class optimum.neuron.NeuronModelForImageClassification

< source >

参数

config (transformers.PretrainedConfig) — PretrainedConfig 是模型配置类，包含模型的所有参数。使用配置文件初始化不会加载与模型关联的权重，只会加载配置。请查看 optimum.neuron.modeling.NeuronTracedModel.from_pretrained 方法来加载模型权重。
model (torch.jit._script.ScriptModule) — torch.jit._script.ScriptModule 是 TorchScript 模块，其中嵌入了由 neuron(x) 编译器编译的 NEFF（神经元可执行文件格式）。

神经元模型，顶部带有一个图像分类头（[CLS] 令牌的最终隐藏状态之上的线性层），例如用于 ImageNet。

此模型继承自 ~neuron.modeling.NeuronTracedModel。查看超类文档，了解库为所有模型实现的通用方法（例如下载或保存）

用于图像分类任务的神经元模型。此类正式支持 beit、convnext、convnextv2、deit、levit、mobilenet_v2、mobilevit、vit 等。

forward

< source >

( pixel_values: Tensor **kwargs )

参数

pixel_values (Union[torch.Tensor, None]，形状为 (batch_size, num_channels, height, width)，默认为 None) — 与当前批次中的图像对应的像素值。像素值可以使用 AutoImageProcessor 从编码图像中获得。

NeuronModelForImageClassification 的 forward 方法，覆盖了 __call__ 特殊方法。

图像分类示例

>>> import requests
>>> from PIL import Image
>>> from optimum.neuron import NeuronModelForImageClassification
>>> from transformers import AutoImageProcessor

>>> url = "http://images.cocodataset.org/val2017/000000039769.jpg"
>>> image = Image.open(requests.get(url, stream=True).raw)

>>> preprocessor = AutoImageProcessor.from_pretrained("optimum/vit-base-patch16-224-neuronx")
>>> model = NeuronModelForImageClassification.from_pretrained("optimum/vit-base-patch16-224-neuronx")

>>> inputs = preprocessor(images=image, return_tensors="pt")

>>> outputs = model(**inputs)
>>> logits = outputs.logits
>>> predicted_label = logits.argmax(-1).item()

使用 optimum.neuron.pipeline 的示例

>>> import requests
>>> from PIL import Image
>>> from transformers import AutoImageProcessor
>>> from optimum.neuron import NeuronModelForImageClassification, pipeline

>>> preprocessor = AutoImageProcessor.from_pretrained("optimum/vit-base-patch16-224-neuronx")
>>> model = NeuronModelForImageClassification.from_pretrained("optimum/vit-base-patch16-224-neuronx")
>>> pipe = pipeline("image-classification", model=model, feature_extractor=preprocessor)

>>> url = "http://images.cocodataset.org/val2017/000000039769.jpg"
>>> pred = pipe(url)

NeuronModelForSemanticSegmentation

class optimum.neuron.NeuronModelForSemanticSegmentation

< source >

参数

config (transformers.PretrainedConfig) — PretrainedConfig 是模型配置类，包含模型的所有参数。使用配置文件初始化不会加载与模型关联的权重，只会加载配置。请查看 optimum.neuron.modeling.NeuronTracedModel.from_pretrained 方法来加载模型权重。
model (torch.jit._script.ScriptModule) — torch.jit._script.ScriptModule 是 TorchScript 模块，其中嵌入了由 neuron(x) 编译器编译的 NEFF（神经元可执行文件格式）。

神经元模型，顶部带有语义分割头，例如用于 Pascal VOC。

此模型继承自 ~neuron.modeling.NeuronTracedModel。查看超类文档，了解库为所有模型实现的通用方法（例如下载或保存）

用于语义分割的神经元模型，顶部带有全 MLP 解码头，例如用于 ADE20k、CityScapes。此类正式支持 mobilevit、mobilenet-v2 等。

forward

< source >

( pixel_values: Tensor **kwargs )

参数

pixel_values (Union[torch.Tensor, None]，形状为 (batch_size, num_channels, height, width)，默认为 None) — 与当前批次中的图像对应的像素值。像素值可以使用 AutoImageProcessor 从编码图像中获得。

NeuronModelForSemanticSegmentation 的 forward 方法，覆盖了 __call__ 特殊方法。

语义分割示例

>>> import requests
>>> from PIL import Image
>>> from optimum.neuron import NeuronModelForSemanticSegmentation
>>> from transformers import AutoImageProcessor

>>> url = "http://images.cocodataset.org/val2017/000000039769.jpg"
>>> image = Image.open(requests.get(url, stream=True).raw)

>>> preprocessor = AutoImageProcessor.from_pretrained("optimum/deeplabv3-mobilevit-small-neuronx")
>>> model = NeuronModelForSemanticSegmentation.from_pretrained("optimum/deeplabv3-mobilevit-small-neuronx")

>>> inputs = preprocessor(images=image, return_tensors="pt")

>>> outputs = model(**inputs)
>>> logits = outputs.logits

使用 optimum.neuron.pipeline 的示例

>>> import requests
>>> from PIL import Image
>>> from transformers import AutoImageProcessor
>>> from optimum.neuron import NeuronModelForSemanticSegmentation, pipeline

>>> preprocessor = AutoImageProcessor.from_pretrained("optimum/deeplabv3-mobilevit-small-neuronx")
>>> model = NeuronModelForSemanticSegmentation.from_pretrained("optimum/deeplabv3-mobilevit-small-neuronx")
>>> pipe = pipeline("image-segmentation", model=model, feature_extractor=preprocessor)

>>> url = "http://images.cocodataset.org/val2017/000000039769.jpg"
>>> pred = pipe(url)

NeuronModelForObjectDetection

class optimum.neuron.NeuronModelForObjectDetection

< source >

参数

config (transformers.PretrainedConfig) — PretrainedConfig 是模型配置类，包含模型的所有参数。使用配置文件初始化不会加载与模型关联的权重，只会加载配置。请查看 optimum.neuron.modeling.NeuronTracedModel.from_pretrained 方法来加载模型权重。
model (torch.jit._script.ScriptModule) — torch.jit._script.ScriptModule 是 TorchScript 模块，其中嵌入了由 neuron(x) 编译器编译的 NEFF（神经元可执行文件格式）。

神经元模型，顶部带有对象检测头，用于诸如 COCO 检测之类的任务。

此模型继承自 ~neuron.modeling.NeuronTracedModel。查看超类文档，了解库为所有模型实现的通用方法（例如下载或保存）

用于对象检测的神经元模型，顶部带有对象检测头，用于诸如 COCO 检测之类的任务。

forward

< source >

( pixel_values: Tensor **kwargs )

参数

pixel_values (Union[torch.Tensor, None]，形状为 (batch_size, num_channels, height, width)，默认为 None) — 与当前批次中的图像对应的像素值。像素值可以使用 AutoImageProcessor 从编码图像中获得。

NeuronModelForObjectDetection 的 forward 方法，覆盖了 __call__ 特殊方法。

对象检测示例

>>> import requests
>>> from PIL import Image
>>> from optimum.neuron import NeuronModelForObjectDetection
>>> from transformers import AutoImageProcessor

>>> url = "http://images.cocodataset.org/val2017/000000039769.jpg"
>>> image = Image.open(requests.get(url, stream=True).raw)

>>> preprocessor = AutoImageProcessor.from_pretrained("hustvl/yolos-tiny")
>>> model = NeuronModelForObjectDetection.from_pretrained("hustvl/yolos-tiny", export=True, batch_size=1)

>>> inputs = preprocessor(images=image, return_tensors="pt")

>>> outputs = model(**inputs)
>>> target_sizes = torch.tensor([image.size[::-1]])
>>> results = image_processor.post_process_object_detection(outputs, threshold=0.9, target_sizes=target_sizes)[0]

使用 optimum.neuron.pipeline 的示例

>>> import requests
>>> from PIL import Image
>>> from transformers import AutoImageProcessor
>>> from optimum.neuron import NeuronModelForObjectDetection, pipeline

>>> preprocessor = AutoImageProcessor.from_pretrained("hustvl/yolos-tiny")
>>> model = NeuronModelForObjectDetection.from_pretrained("hustvl/yolos-tiny")
>>> pipe = pipeline("object-detection", model=model, feature_extractor=preprocessor)

>>> url = "http://images.cocodataset.org/val2017/000000039769.jpg"
>>> pred = pipe(url)

音频

以下自动类可用于以下音频任务。

NeuronModelForAudioClassification

class optimum.neuron.NeuronModelForAudioClassification

< source >

参数

config (transformers.PretrainedConfig) — PretrainedConfig 是模型配置类，包含模型的所有参数。使用配置文件初始化不会加载与模型关联的权重，只会加载配置。请查看 optimum.neuron.modeling.NeuronTracedModel.from_pretrained 方法来加载模型权重。
model (torch.jit._script.ScriptModule) — torch.jit._script.ScriptModule 是 TorchScript 模块，其中嵌入了由 neuron(x) 编译器编译的 NEFF（神经元可执行文件格式）。

带有音频分类头的神经元模型。

此模型继承自 ~neuron.modeling.NeuronTracedModel。查看超类文档，了解库为所有模型实现的通用方法（例如下载或保存）

用于音频分类的神经元模型，顶部带有序列分类头（池化输出之上的线性层），用于诸如 SUPERB 关键词识别之类的任务。

forward

< source >

( input_values: Tensor **kwargs )

参数

input_values (torch.Tensor，形状为 (batch_size, sequence_length)) — 输入原始语音波形的浮点值。输入值可以使用 AutoProcessor 从加载到数组中的音频文件中获得。

NeuronModelForAudioClassification 的 forward 方法，覆盖了 __call__ 特殊方法。

音频分类示例

>>> from transformers import AutoProcessor
>>> from optimum.neuron import NeuronModelForAudioClassification
>>> from datasets import load_dataset
>>> import torch

>>> dataset = load_dataset("hf-internal-testing/librispeech_asr_demo", "clean", split="validation")
>>> dataset = dataset.sort("id")
>>> sampling_rate = dataset.features["audio"].sampling_rate

>>> feature_extractor = AutoProcessor.from_pretrained("Jingya/wav2vec2-large-960h-lv60-self-neuronx-audio-classification")
>>> model = NeuronModelForAudioClassification.from_pretrained("Jingya/wav2vec2-large-960h-lv60-self-neuronx-audio-classification")

>>> # audio file is decoded on the fly
>>> inputs = feature_extractor(dataset[0]["audio"]["array"], sampling_rate=sampling_rate, return_tensors="pt")

>>> logits = model(**inputs).logits
>>> predicted_class_ids = torch.argmax(logits, dim=-1).item()
>>> predicted_label = model.config.id2label[predicted_class_ids]

使用 optimum.neuron.pipeline 的示例

>>> from transformers import AutoProcessor
>>> from optimum.neuron import NeuronModelForAudioClassification, pipeline

>>> feature_extractor = AutoProcessor.from_pretrained("Jingya/wav2vec2-large-960h-lv60-self-neuronx-audio-classification")
>>> dataset = load_dataset("hf-internal-testing/librispeech_asr_demo", "clean", split="validation")
>>> dataset = dataset.sort("id")

>>> model = NeuronModelForAudioClassification.from_pretrained("Jingya/wav2vec2-large-960h-lv60-self-neuronx-audio-classification")
>>> ac = pipeline("audio-classification", model=model, feature_extractor=feature_extractor)

>>> pred = ac(dataset[0]["audio"]["array"])

NeuronModelForAudioFrameClassification

class optimum.neuron.NeuronModelForAudioFrameClassification

< source >

参数

config (transformers.PretrainedConfig) — PretrainedConfig 是模型配置类，包含模型的所有参数。使用配置文件初始化不会加载与模型关联的权重，只会加载配置。请查看 optimum.neuron.modeling.NeuronTracedModel.from_pretrained 方法来加载模型权重。
model (torch.jit._script.ScriptModule) — torch.jit._script.ScriptModule 是由 neuron(x) 编译器编译，嵌入了 NEFF(Neuron Executable File Format) 的 TorchScript 模块。

带有音频帧分类头的 Neuron 模型。

此模型继承自 ~neuron.modeling.NeuronTracedModel。查看超类文档，了解库为所有模型实现的通用方法（例如下载或保存）

带有帧分类头的 Neuron 模型，用于说话人日志等任务。

forward

< source >

( input_values: Tensor **kwargs )

参数

input_values (形状为 (batch_size, sequence_length) 的 torch.Tensor) — 输入原始语音波形的浮点值。输入值可以从加载到数组中的音频文件中获得，使用 AutoProcessor。

NeuronModelForAudioFrameClassification 的 forward 方法，覆盖了 __call__ 特殊方法。

音频帧分类示例

>>> from transformers import AutoProcessor
>>> from optimum.neuron import NeuronModelForAudioFrameClassification
>>> from datasets import load_dataset
>>> import torch

>>> dataset = load_dataset("hf-internal-testing/librispeech_asr_demo", "clean", split="validation")
>>> dataset = dataset.sort("id")
>>> sampling_rate = dataset.features["audio"].sampling_rate

>>> feature_extractor = AutoProcessor.from_pretrained("Jingya/wav2vec2-base-superb-sd-neuronx")
>>> model =  NeuronModelForAudioFrameClassification.from_pretrained("Jingya/wav2vec2-base-superb-sd-neuronx")

>>> inputs = feature_extractor(dataset[0]["audio"]["array"], return_tensors="pt", sampling_rate=sampling_rate)
>>> logits = model(**inputs).logits

>>> probabilities = torch.sigmoid(logits[0])
>>> labels = (probabilities > 0.5).long()
>>> labels[0].tolist()

NeuronModelForCTC

class optimum.neuron.NeuronModelForCTC

< source >

参数

config (transformers.PretrainedConfig) — PretrainedConfig 是模型配置类，包含模型的所有参数。使用配置文件初始化不会加载与模型关联的权重，仅加载配置。查看 optimum.neuron.modeling.NeuronTracedModel.from_pretrained 方法来加载模型权重。
model (torch.jit._script.ScriptModule) — torch.jit._script.ScriptModule 是由 neuron(x) 编译器编译，嵌入了 NEFF(Neuron Executable File Format) 的 TorchScript 模块。

带有连接主义时间分类头的 Neuron 模型。

此模型继承自 ~neuron.modeling.NeuronTracedModel。查看超类文档，了解库为所有模型实现的通用方法（例如下载或保存）

带有用于连接主义时间分类 (CTC) 的语言建模头的 Neuron 模型。

forward

< source >

( input_values: Tensor **kwargs )

参数

input_values (形状为 (batch_size, sequence_length) 的 torch.Tensor) — 输入原始语音波形的浮点值。输入值可以从加载到数组中的音频文件中获得，使用 AutoProcessor。

NeuronModelForCTC 的 forward 方法，覆盖了 __call__ 特殊方法。

CTC 示例

>>> from transformers import AutoProcessor, Wav2Vec2ForCTC
>>> from optimum.neuron import NeuronModelForCTC
>>> from datasets import load_dataset
>>> import torch

>>> dataset = load_dataset("hf-internal-testing/librispeech_asr_demo", "clean", split="validation")
>>> dataset = dataset.sort("id")
>>> sampling_rate = dataset.features["audio"].sampling_rate

>>> processor = AutoProcessor.from_pretrained("Jingya/wav2vec2-large-960h-lv60-self-neuronx-ctc")
>>> model = NeuronModelForCTC.from_pretrained("Jingya/wav2vec2-large-960h-lv60-self-neuronx-ctc")

>>> # audio file is decoded on the fly
>>> inputs = processor(dataset[0]["audio"]["array"], sampling_rate=sampling_rate, return_tensors="pt")
>>> logits = model(**inputs).logits
>>> predicted_ids = torch.argmax(logits, dim=-1)

>>> transcription = processor.batch_decode(predicted_ids)

使用 optimum.neuron.pipeline 的示例

>>> from transformers import AutoProcessor
>>> from optimum.neuron import NeuronModelForCTC, pipeline

>>> processor = AutoProcessor.from_pretrained("Jingya/wav2vec2-large-960h-lv60-self-neuronx-ctc")
>>> dataset = load_dataset("hf-internal-testing/librispeech_asr_demo", "clean", split="validation")
>>> dataset = dataset.sort("id")

>>> model = NeuronModelForCTC.from_pretrained("Jingya/wav2vec2-large-960h-lv60-self-neuronx-ctc")
>>> asr = pipeline("automatic-speech-recognition", model=model, feature_extractor=processor.feature_extractor, tokenizer=processor.tokenizer)

NeuronModelForXVector

class optimum.neuron.NeuronModelForXVector

< source >

参数

config (transformers.PretrainedConfig) — PretrainedConfig 是模型配置类，包含模型的所有参数。使用配置文件初始化不会加载与模型关联的权重，仅加载配置。查看 optimum.neuron.modeling.NeuronTracedModel.from_pretrained 方法来加载模型权重。
model (torch.jit._script.ScriptModule) — torch.jit._script.ScriptModule 是由 neuron(x) 编译器编译，嵌入了 NEFF(Neuron Executable File Format) 的 TorchScript 模块。

带有用于说话人验证等任务的 XVector 特征提取头的 Neuron 模型。

此模型继承自 ~neuron.modeling.NeuronTracedModel。查看超类文档，了解库为所有模型实现的通用方法（例如下载或保存）

带有用于说话人验证等任务的 XVector 特征提取头的 Neuron 模型。

forward

< source >

( input_values: Tensor **kwargs )

参数

input_values (形状为 (batch_size, sequence_length) 的 torch.Tensor) — 输入原始语音波形的浮点值。输入值可以从加载到数组中的音频文件中获得，使用 AutoProcessor。

NeuronModelForXVector 的 forward 方法，覆盖了 __call__ 特殊方法。

音频 XVector 示例

>>> from transformers import AutoProcessor
>>> from optimum.neuron import NeuronModelForXVector
>>> from datasets import load_dataset
>>> import torch

>>> dataset = load_dataset("hf-internal-testing/librispeech_asr_demo", "clean", split="validation")
>>> dataset = dataset.sort("id")
>>> sampling_rate = dataset.features["audio"].sampling_rate

>>> feature_extractor = AutoProcessor.from_pretrained("Jingya/wav2vec2-base-superb-sv-neuronx")
>>> model = NeuronModelForXVector.from_pretrained("Jingya/wav2vec2-base-superb-sv-neuronx")

>>> inputs = feature_extractor(
...     [d["array"] for d in dataset[:2]["audio"]], sampling_rate=sampling_rate, return_tensors="pt", padding=True
... )
>>> embeddings = model(**inputs).embeddings

>>> embeddings = torch.nn.functional.normalize(embeddings, dim=-1)

>>> cosine_sim = torch.nn.CosineSimilarity(dim=-1)
>>> similarity = cosine_sim(embeddings[0], embeddings[1])
>>> threshold = 0.7
>>> if similarity < threshold:
...     print("Speakers are not the same!")
>>> round(similarity.item(), 2)

NeuronWhisperForConditionalGeneration

class optimum.neuron.NeuronWhisperForConditionalGeneration

< source >

参数

encoder (torch.jit._script.ScriptModule) — torch.jit._script.ScriptModule 是编码器的 TorchScript 模块，它由 neuron(x) 编译器编译，并嵌入了 NEFF(Neuron Executable File Format)。
decoder (torch.jit._script.ScriptModule) — torch.jit._script.ScriptModule 是解码器的 TorchScript 模块，它由 neuron(x) 编译器编译，并嵌入了 NEFF(Neuron Executable File Format)。
config (transformers.PretrainedConfig) — PretrainedConfig 是模型配置类，包含模型的所有参数。使用配置文件初始化不会加载与模型关联的权重，仅加载配置。查看 optimum.neuron.modeling.NeuronTracedModel.from_pretrained 方法来加载模型权重。

带有语言建模头的 Whisper Neuron 模型，可用于自动语音识别。

此模型继承自 ~neuron.modeling.NeuronTracedModel。查看超类文档，了解库为所有模型实现的通用方法（例如下载或保存）

forward

< source >

( input_features: typing.Optional[torch.FloatTensor] = None decoder_input_ids: typing.Optional[torch.LongTensor] = None encoder_outputs: typing.Optional[typing.Tuple[torch.FloatTensor]] = None **kwargs )

参数

input_features (形状为 (batch_size, feature_size, sequence_length) 的 Optional[torch.FloatTensor]) — 从原始语音波形中提取的浮点值梅尔特征。原始语音波形可以通过将 .flac 或 .wav 音频文件加载到 List[float] 类型或 numpy.ndarray 类型的数组中获得，例如通过 soundfile 库 (pip install soundfile)。为了将数组准备成 input_features，应该使用 AutoFeatureExtractor 来提取梅尔特征、进行填充并转换为 torch.FloatTensor 类型的张量。参见 ~WhisperFeatureExtractor.__call__
decoder_input_ids (形状为 (batch_size, max_sequence_length) 的 Optional[torch.LongTensor]) — 词汇表中解码器输入序列令牌的索引。索引可以使用 WhisperTokenizer 获得。有关详细信息，请参见 PreTrainedTokenizer.encode 和 PreTrainedTokenizer.__call__。由于 Whisper 尚不支持缓存，因此需要将其填充到用于编译的 sequence_length。
encoder_outputs (Optional[Tuple[torch.FloatTensor]]) — 元组由形状为 (batch_size, sequence_length, hidden_size) 的 last_hidden_state 组成，它是编码器最后一层输出的隐藏状态序列。在解码器的交叉注意力中使用。

NeuronWhisperForConditionalGeneration 的 forward 方法，覆盖了 __call__ 特殊方法。

（以下模型使用 neuronx 编译器编译，只能在 INF2 上运行。）

使用 Whisper 模型进行自动语音识别的示例

from datasets import load_dataset
from transformers import AutoProcessor
from optimum.neuron import NeuronWhisperForConditionalGeneration

# Select an audio file and read it:
ds = load_dataset("hf-internal-testing/librispeech_asr_dummy", "clean", split="validation")
audio_sample = ds[1]["audio"]
processor = AutoProcessor.from_pretrained(openai/whisper-tiny)

# Use the model and processor to transcribe the audio:
input_features = processor(
    audio_sample["array"], sampling_rate=audio_sample["sampling_rate"], return_tensors="pt"
).input_features

# Compile the model to Neuron format
neuron_model = NeuronWhisperForConditionalGeneration.from_pretrained(openai/whisper-tiny, export=True, batch_size=1, sequence_length=128)
neuron_model.save_pretrained("whisper_tiny_neuronx/")
del neuron_model

# Inference
neuron_model = NeuronWhisperForConditionalGeneration.from_pretrained("whisper_tiny_neuronx/")
predicted_ids = neuron_model.generate(input_features)
transcription = processor.batch_decode(predicted_ids, skip_special_tokens=True)

Stable Diffusion

以下 Neuron 模型类可用于 stable diffusion 任务。

NeuronStableDiffusionPipeline

class optimum.neuron.NeuronStableDiffusionControlNetPipeline
< source >
( config: typing.Dict[str, typing.Any] configs: typing.Dict[str, ForwardRef('PretrainedConfig')] neuron_configs: typing.Dict[str, ForwardRef('NeuronDefaultConfig')] data_parallel_mode: typing.Literal['none', 'unet', 'transformer', 'all'] scheduler: typing.Optional[diffusers.schedulers.scheduling_utils.SchedulerMixin] vae_decoder: typing.Union[torch.jit._script.ScriptModule, ForwardRef('NeuronModelVaeDecoder')] text_encoder: typing.Union[torch.jit._script.ScriptModule, ForwardRef('NeuronModelTextEncoder'), NoneType] = None text_encoder_2: typing.Union[torch.jit._script.ScriptModule, ForwardRef('NeuronModelTextEncoder'), NoneType] = None unet: typing.Union[torch.jit._script.ScriptModule, ForwardRef('NeuronModelUnet'), NoneType] = None transformer: typing.Union[torch.jit._script.ScriptModule, ForwardRef('NeuronModelTransformer'), NoneType] = None vae_encoder: typing.Union[torch.jit._script.ScriptModule, ForwardRef('NeuronModelVaeEncoder'), NoneType] = None image_encoder: typing.Optional[torch.jit._script.ScriptModule] = None safety_checker: typing.Optional[torch.jit._script.ScriptModule] = None tokenizer: typing.Union[transformers.models.clip.tokenization_clip.CLIPTokenizer, transformers.utils.dummy_sentencepiece_objects.T5Tokenizer, NoneType] = None tokenizer_2: typing.Optional[transformers.models.clip.tokenization_clip.CLIPTokenizer] = None feature_extractor: typing.Optional[transformers.models.clip.feature_extraction_clip.CLIPFeatureExtractor] = None controlnet: typing.Union[torch.jit._script.ScriptModule, typing.List[torch.jit._script.ScriptModule], ForwardRef('NeuronControlNetModel'), ForwardRef('NeuronMultiControlNetModel'), NoneType] = None requires_aesthetics_score: bool = False force_zeros_for_empty_prompt: bool = True add_watermarker: typing.Optional[bool] = None model_save_dir: typing.Union[str, pathlib.Path, tempfile.TemporaryDirectory, NoneType] = None model_and_config_save_paths: typing.Optional[typing.Dict[str, typing.Tuple[str, pathlib.Path]]] = None )

call
< source >
( prompt: typing.Union[str, typing.List[str]] = None image: typing.Union[PIL.Image.Image, numpy.ndarray, torch.Tensor, typing.List[PIL.Image.Image], typing.List[numpy.ndarray], typing.List[torch.Tensor]] = None num_inference_steps: int = 50 timesteps: typing.Optional[typing.List[int]] = None sigmas: typing.Optional[typing.List[float]] = None guidance_scale: float = 7.5 negative_prompt: typing.Union[str, typing.List[str], NoneType] = None num_images_per_prompt: typing.Optional[int] = 1 eta: float = 0.0 generator: typing.Union[torch._C.Generator, typing.List[torch._C.Generator], NoneType] = None latents: typing.Optional[torch.Tensor] = None prompt_embeds: typing.Optional[torch.Tensor] = None negative_prompt_embeds: typing.Optional[torch.Tensor] = None ip_adapter_image: typing.Union[PIL.Image.Image, numpy.ndarray, torch.Tensor, typing.List[PIL.Image.Image], typing.List[numpy.ndarray], typing.List[torch.Tensor], NoneType] = None ip_adapter_image_embeds: typing.Optional[typing.List[torch.Tensor]] = None output_type: str = 'pil' return_dict: bool = True cross_attention_kwargs: typing.Optional[typing.Dict[str, typing.Any]] = None controlnet_conditioning_scale: typing.Union[float, typing.List[float]] = 1.0 guess_mode: bool = False control_guidance_start: typing.Union[float, typing.List[float]] = 0.0 control_guidance_end: typing.Union[float, typing.List[float]] = 1.0 clip_skip: typing.Optional[int] = None callback_on_step_end: typing.Union[typing.Callable[[int, int, typing.Dict], NoneType], diffusers.callbacks.PipelineCallback, diffusers.callbacks.MultiPipelineCallbacks, NoneType] = None callback_on_step_end_tensor_inputs: typing.List[str] = ['latents'] kwargs ) → `diffusers.pipelines.stable_diffusion.StableDiffusionPipelineOutput` or `tuple`

参数

prompt (`Optional[Union[str, List[str]]]`, defaults to `None`) — The prompt or prompts to guide image generation. If not defined, you need to pass `prompt_embeds`.
image (`Optional["PipelineImageInput"]`, defaults to `None`) — The ControlNet input condition to provide guidance to the `unet` for generation. If the type is specified as `torch.Tensor`, it is passed to ControlNet as is. `PIL.Image.Image` can also be accepted as an image. The dimensions of the output image defaults to `image`’s dimensions. If height and/or width are passed, `image` is resized accordingly. If multiple ControlNets are specified in `init`, images must be passed as a list such that each element of the list can be correctly batched for input to a single ControlNet. When `prompt` is a list, and if a list of images is passed for a single ControlNet, each will be paired with each prompt in the `prompt` list. This also applies to multiple ControlNets, where a list of image lists can be passed to batch for each prompt and each ControlNet.
num_inference_steps (`int`, defaults to 50) — The number of denoising steps. More denoising steps usually lead to a higher quality image at the expense of slower inference.
timesteps (`Optional[List[int]]`, defaults to `None`) — Custom timesteps to use for the denoising process with schedulers which support a `timesteps` argument in their `set_timesteps` method. If not defined, the default behavior when `num_inference_steps` is passed will be used. Must be in descending order.
sigmas (`Optional[List[int]]`, defaults to `None`) — Custom sigmas to use for the denoising process with schedulers which support a `sigmas` argument in their `set_timesteps` method. If not defined, the default behavior when `num_inference_steps` is passed will be used.
guidance_scale (`float`, defaults to 7.5) — A higher guidance scale value encourages the model to generate images closely linked to the text `prompt` at the expense of lower image quality. Guidance scale is enabled when `guidance_scale > 1`.
negative_prompt (`Optional[Union[str, List[str]]]`, defaults to `None`) — The prompt or prompts to guide what to not include in image generation. If not defined, you need to pass `negative_prompt_embeds` instead. Ignored when not using guidance (`guidance_scale < 1`).
num_images_per_prompt (`int`, defaults to 1) — The number of images to generate per prompt. If it is different from the batch size used for the compiltaion, it will be overriden by the static batch size of neuron (except for dynamic batching).
eta (`float`, defaults to 0.0) — 对应于 DDIM 论文中的参数 eta (η)。仅适用于 `diffusers.schedulers.DDIMScheduler`，在其他调度器中将被忽略。
generator (`Optional[Union[torch.Generator, List[torch.Generator]]]`, defaults to `None`) — 用于使生成结果确定性的 `torch.Generator`。
latents (`Optional[torch.Tensor]`, defaults to `None`) — 预生成的、从高斯分布中采样的噪声潜变量，用作图像生成的输入。可用于通过不同的提示调整相同的生成结果。如果未提供，则会通过使用提供的随机 `generator` 进行采样来生成潜变量张量。
prompt_embeds (`Optional[torch.Tensor]`, defaults to `None`) — 预生成的文本嵌入。可用于轻松调整文本输入（提示权重）。如果未提供，则会从 `prompt` 输入参数生成文本嵌入。
negative_prompt_embeds (`Optional[torch.Tensor]`, defaults to `None`) — 预生成的负面文本嵌入。可用于轻松调整文本输入（提示权重）。如果未提供，则会从 `negative_prompt` 输入参数生成 `negative_prompt_embeds`。
ip_adapter_image — (`Optional[PipelineImageInput]`, defaults to `None`): 与 IP 适配器一起使用的可选图像输入。
ip_adapter_image_embeds (`Optional[List[torch.Tensor]]`, defaults to `None`) — IP 适配器的预生成图像嵌入。它应该是一个列表，其长度与 IP 适配器的数量相同。每个元素都应该是一个形状为 `(batch_size, num_images, emb_dim)` 的张量。如果 `do_classifier_free_guidance` 设置为 `True`，则应包含负面图像嵌入。如果未提供，则从 `ip_adapter_image` 输入参数计算嵌入。
output_type (`str`, defaults to `"pil"`) — 生成图像的输出格式。在 `PIL.Image` 或 `np.array` 之间选择。
return_dict (`bool`, defaults to `True`) — 是否返回 `diffusers.pipelines.stable_diffusion.StableDiffusionPipelineOutput` 而不是普通元组。
cross_attention_kwargs (`Optional[Dict[str, Any]]`, defaults to `None`) — 一个 kwargs 字典，如果指定，则会传递给 `self.processor` 中定义的 `AttentionProcessor`。
controlnet_conditioning_scale (`Union[float, List[float]]`, defaults to 1.0) — ControlNet 的输出在添加到原始 `unet` 中的残差之前，会乘以 `controlnet_conditioning_scale`。如果在 `init` 中指定了多个 ControlNet，则可以将相应的比例设置为列表。
guess_mode (`bool`, defaults to `False`) — 即使您删除了所有提示，ControlNet 编码器也会尝试识别输入图像的内容。建议 `guidance_scale` 值在 3.0 到 5.0 之间。
control_guidance_start (`Union[float, List[float]]`, defaults to 0.0) — ControlNet 开始应用的占总步数的百分比。
control_guidance_end** (`Union[float, List[float]]`, optional, defaults to 1.0) — ControlNet 停止应用的占总步数的百分比。
clip_skip (`Optional[int]`, defaults to `None`) — 计算提示嵌入时要从 CLIP 中跳过的层数。值为 1 表示预最终层的输出将用于计算提示嵌入。
callback_on_step_end (`Optional[Union[Callable[[int, int, Dict], None], PipelineCallback, MultiPipelineCallbacks]]`, defaults to `None`) — 一个函数或 `PipelineCallback` 或 `MultiPipelineCallbacks` 的子类，它在推理期间的每个去噪步骤结束时被调用。具有以下参数：`callback_on_step_end(self: DiffusionPipeline, step: int, timestep: int, callback_kwargs: Dict)`。 `callback_kwargs` 将包括 `callback_on_step_end_tensor_inputs` 指定的所有张量列表。
callback_on_step_end_tensor_inputs (`List[str]`, defaults to `["latents"]`) — `callback_on_step_end` 函数的张量输入列表。列表中指定的张量将作为 `callback_kwargs` 参数传递。您将只能包含管道类 `._callback_tensor_inputs` 属性中列出的变量。

返回

`diffusers.pipelines.stable_diffusion.StableDiffusionPipelineOutput` 或 `tuple`

如果 `return_dict` 为 `True`，则返回 `diffusers.pipelines.stable_diffusion.StableDiffusionPipelineOutput`，否则返回 `tuple`，其中第一个元素是包含生成图像的列表，第二个元素是 `bool` 列表，指示相应的生成图像是否包含“不适合工作场所观看”（nsfw）内容。

调用管道进行生成的功能。

NeuronPixArtAlphaPipeline

参数

prompt (Optional[Union[str, List[str]]], defaults to None) — 用于引导图像生成的提示或提示语。如果未定义，则需要传递 prompt_embeds。
prompt_2 (Optional[Union[str, List[str]]], defaults to None) — 要发送到 tokenizer_2 和 text_encoder_2 的提示或提示语。如果未定义，则 prompt 将在两个文本编码器中使用。
image (Optional["PipelineImageInput"], defaults to None) — ControlNet 输入条件，为 unet 的生成提供指导。如果类型指定为 torch.Tensor，则会按原样传递给 ControlNet。PIL.Image.Image 也可以接受作为图像。输出图像的尺寸默认为 image 的尺寸。如果传递了高度和/或宽度，则会相应地调整 image 的大小。如果在 init 中指定了多个 ControlNet，则必须将图像作为列表传递，以便列表的每个元素都可以正确地批量处理，以输入到单个 ControlNet。
num_inference_steps (int, defaults to 50) — 去噪步骤的数量。更多的去噪步骤通常会带来更高质量的图像，但代价是推理速度较慢。
timesteps (Optional[List[int]], defaults to None) — 用于去噪过程的自定义时间步长，适用于调度器，这些调度器在其 set_timesteps 方法中支持 timesteps 参数。如果未定义，则将使用传递 num_inference_steps 时的默认行为。必须按降序排列。
sigmas (Optional[List[int]], defaults to None) — 用于去噪过程的自定义 sigmas，适用于调度器，这些调度器在其 set_timesteps 方法中支持 sigmas 参数。如果未定义，则将使用传递 num_inference_steps 时的默认行为。
denoising_end (Optional[float], defaults to None) — 指定后，确定在有意提前终止之前要完成的总去噪过程的分数（介于 0.0 和 1.0 之间）。因此，返回的样本仍将保留大量噪声，这由调度器选择的离散时间步长决定。当此管道构成“去噪器混合”多管道设置的一部分时，应理想地利用 denoising_end 参数，如 优化图像输出 中详述的那样。
guidance_scale (float, defaults to 5.0) — 更高的 guidance scale 值会鼓励模型生成与文本 prompt 紧密相关的图像，但会牺牲图像质量。当 guidance_scale > 1 时，guidance scale 生效。
negative_prompt (Optional[Union[str, List[str]]], defaults to None) — 用于引导图像生成中不应包含的内容的提示或提示语。如果未定义，则需要改为传递 negative_prompt_embeds。当不使用 guidance 时（guidance_scale < 1），将被忽略。
negative_prompt_2 (Optional[Union[str, List[str]]], defaults to None) — 用于引导图像生成中不应包含的内容的提示或提示语。这会发送到 tokenizer_2 和 text_encoder_2。如果未定义，则 negative_prompt 将在两个文本编码器中使用。
num_images_per_prompt (int, defaults to 1) — 每个 prompt 生成的图像数量。
eta (float, defaults to 0.0) — 对应于 DDIM 论文中的参数 eta (η)。仅适用于 diffusers.schedulers.DDIMScheduler，在其他调度器中将被忽略。
generator (Optional[Union[torch.Generator, List[torch.Generator]]], defaults to None) — 用于使生成具有确定性的 torch.Generator。
latents (Optional[torch.Tensor], defaults to None) — 预生成的、从高斯分布中采样的噪声 latents，用作图像生成的输入。可用于使用不同的 prompt 微调相同的生成结果。如果未提供，则会通过使用提供的随机 generator 进行采样来生成 latents 张量。
prompt_embeds (Optional[torch.Tensor], 默认为 None) — 预生成的文本嵌入。可用于轻松调整文本输入（提示权重）。如果未提供，则会从 prompt 输入参数生成文本嵌入。
negative_prompt_embeds (Optional[torch.Tensor], 默认为 None) — 预生成的负面文本嵌入。可用于轻松调整文本输入（提示权重）。如果未提供，则会从 negative_prompt 输入参数生成 negative_prompt_embeds。
pooled_prompt_embeds (Optional[torch.Tensor], 默认为 None) — 预生成的池化文本嵌入。可用于轻松调整文本输入（提示权重）。如果未提供，则会从 prompt 输入参数生成池化文本嵌入。
negative_pooled_prompt_embeds (Optional[torch.Tensor], 默认为 None) — 预生成的负面池化文本嵌入。可用于轻松调整文本输入（提示权重）。如果未提供，则会从 negative_prompt 输入参数生成池化的 negative_prompt_embeds。
ip_adapter_image — (Optional[PipelineImageInput], 默认为 None): 用于 IP 适配器的可选图像输入。
ip_adapter_image_embeds (Optional[List[torch.Tensor]], 默认为 None) — IP 适配器的预生成图像嵌入。它应该是一个列表，长度与 IP 适配器的数量相同。每个元素都应该是一个形状为 (batch_size, num_images, emb_dim) 的张量。如果 do_classifier_free_guidance 设置为 True，则应包含负面图像嵌入。如果未提供，则会从 ip_adapter_image 输入参数计算嵌入。
output_type (Optional[str], 默认为 "pil") — 生成图像的输出格式。在 PIL.Image 或 np.array 之间选择。
return_dict (bool, 默认为 True) — 是否返回 ~pipelines.stable_diffusion.StableDiffusionPipelineOutput 而不是纯元组。
cross_attention_kwargs (Optional[Dict[str, Any]], 默认为 None) — 一个 kwargs 字典，如果指定，则会传递给 self.processor 中定义的 AttentionProcessor。
controlnet_conditioning_scale (Union[float, List[float]], 默认为 1.0) — ControlNet 的输出在添加到原始 unet 中的残差之前，会乘以 controlnet_conditioning_scale。如果在 init 中指定了多个 ControlNet，则可以将相应的比例设置为列表。
guess_mode (bool, 默认为 False) — 即使您移除所有提示，ControlNet 编码器也会尝试识别输入图像的内容。建议 guidance_scale 值介于 3.0 和 5.0 之间。
control_guidance_start (Union[float, List[float]], 默认为 0.0) — ControlNet 开始应用的占总步数的百分比。
control_guidance_end (Union[float, List[float]], 默认为 1.0) — ControlNet 停止应用的占总步数的百分比。
original_size (Optional[Tuple[int, int]], 默认为 (1024, 1024)) — 如果 original_size 与 target_size 不同，则图像将显示为缩小或放大。如果未指定，original_size 默认为 (height, width)。SDXL 微调的一部分，如 https://huggingface.ac.cn/papers/2307.01952 的第 2.2 节中所述。
crops_coords_top_left (Tuple[int, int], 默认为 (0, 0)) — crops_coords_top_left 可用于生成从 crops_coords_top_left 位置向下“裁剪”的图像。通常通过将 crops_coords_top_left 设置为 (0, 0) 来获得良好居中的图像。SDXL 微调的一部分，如 https://huggingface.ac.cn/papers/2307.01952 的第 2.2 节中所述。
target_size (Optional[Tuple[int, int]], 默认为 None) — 在大多数情况下，target_size 应设置为生成图像所需的 height 和 width。如果未指定，它将默认为 (height, width)。SDXL 微调的一部分，如 https://huggingface.ac.cn/papers/2307.01952 的第 2.2 节中所述。
negative_original_size (Optional[Tuple[int, int]], 默认为 None) — 为了基于特定的图像分辨率对生成过程进行负面调节。SDXL 微调的一部分，如 https://huggingface.ac.cn/papers/2307.01952 的第 2.2 节中所述。有关更多信息，请参阅此问题线程：https://github.com/huggingface/diffusers/issues/4208。
negative_crops_coords_top_left (Tuple[int, int], 默认为 (0, 0)) — 为了基于特定的裁剪坐标对生成过程进行负面调节。SDXL 微调的一部分，如 https://huggingface.ac.cn/papers/2307.01952 的第 2.2 节中所述。有关更多信息，请参阅此问题线程：https://github.com/huggingface/diffusers/issues/4208。
negative_target_size (Optional[Tuple[int, int]], 默认为 None) — 为了基于目标图像分辨率对生成过程进行负面调节。在大多数情况下，它应该与 target_size 相同。SDXL 微调的一部分，如 https://huggingface.ac.cn/papers/2307.01952 的第 2.2 节中所述。有关更多信息，请参阅此问题线程：https://github.com/huggingface/diffusers/issues/4208。
clip_skip (Optional[int], 默认为 None) — 从 CLIP 中跳过的层数，以计算提示嵌入。值为 1 表示预最终层的输出将用于计算提示嵌入。
callback_on_step_end (Optional[Union[Callable[[int, int, Dict], None], PipelineCallback, MultiPipelineCallbacks]], 默认为 None) — 在推理期间的每个去噪步骤结束时调用的函数或 PipelineCallback 或 MultiPipelineCallbacks 的子类。具有以下参数：callback_on_step_end(self: DiffusionPipeline, step: int, timestep: int, callback_kwargs: Dict)。callback_kwargs 将包含 callback_on_step_end_tensor_inputs 指定的所有张量列表。
callback_on_step_end_tensor_inputs (List[str], 默认为 ["latents"]) — callback_on_step_end 函数的张量输入列表。列表中指定的张量将作为 callback_kwargs 参数传递。您只能包含管道类的 ._callback_tensor_inputs 属性中列出的变量。

diffusers.pipelines.stable_diffusion.StableDiffusionPipelineOutput 或 tuple

如果 return_dict 为 True，则返回 diffusers.pipelines.stable_diffusion.StableDiffusionPipelineOutput，否则返回包含输出图像的 tuple。

调用管道进行生成的功能。

示例

←Neuron 导出器