数据集工具

用于获取样本和不同数据集的前向循环函数的实用函数。

函数

create_forward_loop	创建并返回一个为特定模型、数据集和分词器配置的前向循环函数。
get_dataset_dataloader	获取一个带有目标模型的数据集名称和分词器的数据加载器。

create_forward_loop(model=None, dataset_name='cnn_dailymail', tokenizer=None, batch_size=1, num_samples=512, max_sample_length=512, device=None, include_labels=False, dataloader=None)

创建并返回一个为特定模型、数据集和分词器配置的前向循环函数。

此函数初始化一个前向循环函数，专门用于使用给定的模型和分词器处理指定数据集中的数据批次。当前向循环函数被调用时，它会遍历数据集，应用分词器准备输入数据，将其输入模型，并返回模型的预测结果。

参数: - model: 用于推理的PyTorch模型。 - dataset_name: 要使用的数据集的名称。 - tokenizer: 用于将文本数据预处理为适合模型格式的分词器。 - batch_size: 返回的数据加载器的批量大小。如果提供0，则自动确定batch_size。 - num_samples: 数据集中的样本数量。 - max_sample_length: 样本的最大长度。 - device: 返回的数据加载器的目标设备。 - include_labels: 是否在数据加载器中包含标签。 - dataloader: 如果提供，则使用提供的数据加载器。

量化示例用法：

import modelopt.torch.quantization as mtq

# Initialize model and tokenizer
# ...

# Create forward loop for calibration
forward_loop = create_forward_loop(
    model=model, dataset_name="cnn_dailymail", tokenizer=tokenizer
)

# Quantize the model with the calibration dataset
mtq.quantize(model, quant_cfg, forward_loop=forward_loop)

返回： - function: 一个无需参数即可调用的前向循环函数。当调用时，此函数会遍历由dataset_name指定的数据集。

Parameters:

• 模型 (模块) –

• dataset_name (str) –

• tokenizer (PreTrainedTokenizer | PreTrainedTokenizerFast) –

• batch_size (int) –

• num_samples (int) –

• max_sample_length (int) –

• 设备 (str | 无) –

• include_labels (bool) –

• dataloader (DataLoader) –

Return type:

可调用的

get_dataset_dataloader(dataset_name='cnn_dailymail', tokenizer=None, batch_size=1, num_samples=512, max_sample_length=512, device=None, include_labels=False)

获取一个带有目标模型的数据集名称和分词器的数据加载器。

Parameters:

• dataset_name (str) – 要加载的数据集的名称。

• tokenizer (PreTrainedTokenizer | PreTrainedTokenizerFast) – Hugginface 分词器的实例。

• batch_size (int) – 返回的数据加载器的批量大小。

• num_samples (int) – 数据集中的样本数量。

• max_sample_length (int) – 样本的最大长度。

• device (str | None) – 返回的数据加载器的目标设备。

• include_labels (bool) – 是否在数据加载器中包含标签。

Returns:

一个dataloader的实例。

Return type:

数据加载器