FP16 优化 ONNX 模型

该模型是 cardiffnlp/twitter-xlm-roberta-base-sentiment 的 ONNX-FP16 优化版本。它仅在 GPU 上运行。根据模型不同，ONNX-FP16 版本可能比基础 PyTorch 模型快 2-3 倍。

有关 ONNX-FP16 与 ONNX 及 PyTorch 的性能基准对比，以及用于生成和验证该模型准确性的脚本，请参阅 https://github.com/joaopn/encoder-optimization-guide。

在包含 10000 条 Reddit 评论的测试集上，该模型与 FP32 模型之间的标签概率差异如下：

Mean: 0.00075560
Std Dev: 0.00073272
Min: 0.00000095
Max: 0.01000583
Median: 0.00054353

Quantiles:
  25th percentile: 0.00024071
  50th percentile: 0.00054353
  75th percentile: 0.00103680
  90th percentile: 0.00168392
  95th percentile: 0.00217985
  99th percentile: 0.00333904

使用方法

该模型的生成方式如下：

from optimum.onnxruntime import ORTOptimizer, ORTModelForSequenceClassification, AutoOptimizationConfig
from transformers import AutoTokenizer

model_id = \"cardiffnlp/twitter-xlm-roberta-base-sentiment\"
save_dir = \"./model-onnx-fp16\"

# 1. Export the base model to ONNX
model = ORTModelForSequenceClassification.from_pretrained(model_id, export=True)
tokenizer = AutoTokenizer.from_pretrained(model_id)

# 2. Setup the Optimizer
optimizer = ORTOptimizer.from_pretrained(model)

# 3. Apply O4 Optimization (GPU-only FP16)
optimization_config = AutoOptimizationConfig.O4()

optimizer.optimize(
    save_dir=save_dir,
    optimization_config=optimization_config
)

# 4. Save tokenizer for a complete package
tokenizer.save_pretrained(save_dir)

你需要安装 GPU 版本的 ONNX Runtime。可以通过以下命令安装：

pip install optimum[onnxruntime-gpu] --extra-index-url https://aiinfra.pkgs.visualstudio.com/PublicPackages/_packaging/onnxruntime-cuda-12/pypi/simple/

为方便起见，你可以使用这个 environment.yml 文件来创建包含所有依赖项的 conda 环境。以下是一个优化的批量使用示例：

import pandas as pd
import torch
from tqdm import tqdm
from transformers import AutoTokenizer
from optimum.onnxruntime import ORTModelForSequenceClassification

def sentiment_analysis_batched(df, batch_size, field_name):
    # Replace with your HuggingFace username/model_id after uploading
    model_id = 'YOUR_USERNAME/YOUR_MODEL_ID'
    file_name = 'model.onnx'
    gpu_id = 0
    
    model = ORTModelForSequenceClassification.from_pretrained(model_id, file_name=file_name, provider=\"CUDAExecutionProvider\", provider_options={'device_id': gpu_id})
    device = torch.device(f\"cuda:{gpu_id}\")

    tokenizer = AutoTokenizer.from_pretrained(model_id)

    results = []

    # Precompute id2label mapping
    id2label = model.config.id2label

    total_samples = len(df)
    with tqdm(total=total_samples, desc=\"Processing samples\") as pbar:
        for start_idx in range(0, total_samples, batch_size):
            end_idx = start_idx + batch_size
            texts = df[field_name].iloc[start_idx:end_idx].tolist()

            inputs = tokenizer(texts, padding=True, truncation=True, return_tensors=\"pt\", max_length=512)
            input_ids = inputs['input_ids'].to(device)
            attention_mask = inputs['attention_mask'].to(device)

            with torch.no_grad():
                outputs = model(input_ids, attention_mask=attention_mask)
            predictions = torch.sigmoid(outputs.logits)  # Use sigmoid for multi-label classification

            # Collect predictions on GPU
            results.append(predictions)

            pbar.update(end_idx - start_idx)

    # Concatenate all results on GPU
    all_predictions = torch.cat(results, dim=0).cpu().numpy()

    # Convert to DataFrame
    predictions_df = pd.DataFrame(all_predictions, columns=[id2label[i] for i in range(all_predictions.shape[1])])

    # Add prediction columns to the original DataFrame
    combined_df = pd.concat([df.reset_index(drop=True), predictions_df], axis=1)

    return combined_df

df = pd.read_csv('https://github.com/joaopn/gpu_benchmark_goemotions/raw/main/data/random_sample_10k.csv.gz')
df = sentiment_analysis_batched(df, batch_size=8, field_name='body')