深度学习-PyTorch：02-基于BERT-base打造AI芯片高效 - 简易版压力测试工具

简介：随着人工智能的迅猛发展，AI芯片的性能测试成为了行业关注的焦点。BERT-base，作为一种高级的语言表示模型，因其复杂度和计算需求，可以视为进行高功耗压力测试的理想选择之一。

历史攻略：

深度学习：win10安装PyTorch

深度学习-PyTorch：01-使用gpu进行运算范例

安装依赖：由于网络问题，建议使用国内源

pip install -i https://mirrors.aliyun.com/pypi/simple/ transformers datasets torch torchvision torchaudio

案例源码：网络正常版，如果运行抛出网络原因，请使用本地版。

import torch
from transformers import BertTokenizer, BertForSequenceClassification, AdamW
from torch.utils.data import DataLoader
from datasets import load_dataset

# 加载模型和数据集
model = BertForSequenceClassification.from_pretrained('bert-base-uncased')
tokenizer = BertTokenizer.from_pretrained('bert-base-uncased')
dataset = load_dataset('glue', 'mrpc', split='train')

# 数据预处理
def encode(examples):
    return tokenizer(examples['sentence1'], examples['sentence2'], truncation=True, padding='max_length', max_length=128)

dataset = dataset.map(encode)
dataset.set_format(type='torch', columns=['input_ids', 'attention_mask', 'label'])

# 数据加载器
train_loader = DataLoader(dataset, batch_size=8)

# 选择设备
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model.to(device)

# 创建优化器
optimizer = AdamW(model.parameters(), lr=1e-5)

# 训练模型
for epoch in range(3):  # 多次循环以增加负载
    for batch in train_loader:
        inputs = {k: v.to(device) for k, v in batch.items()}
        outputs = model(**inputs)

        # 计算损失并进行反向传播
        loss = outputs.loss
        loss.backward()
        optimizer.step()
        optimizer.zero_grad()

离线版本，下载模型：

https://huggingface.co/bert-base-uncased/tree/main

# 将下面三个文件下载到本地
config.json
pytorch_model.bin
vocab.txt

离线版本，下载数据集：

https://www.microsoft.com/en-us/download/details.aspx?id=52398

下载到本地后安装

案例源码：本地离线版本运行

# -*- coding: utf-8 -*-
# time: 2024/01/03 12:39
# file: power_stress.py
# 公众号: 玩转测试开发
import torch
from transformers import BertTokenizer, BertForSequenceClassification, AdamW
from torch.utils.data import DataLoader, Dataset


class MRPCDataset(Dataset):
    def __init__(self, tokenizer, file_path):
        self.tokenizer = tokenizer
        self.sentences = []
        self.labels = []

        with open(file_path, encoding="utf8") as f:
            lines = f.readlines()[1:]
            for line in lines:
                parts = line.strip().split('\t')
                if len(parts) == 5:
                    self.labels.append(int(parts[0]))
                    self.sentences.append((parts[3], parts[4]))

    def __len__(self):
        return len(self.labels)

    def __getitem__(self, idx):
        sentence1, sentence2 = self.sentences[idx]
        encoding = self.tokenizer(sentence1, sentence2, return_tensors='pt',
                                  padding='max_length', truncation=True, max_length=128)
        return {'input_ids': encoding['input_ids'].squeeze(0),  # 删除批处理维度
                'attention_mask': encoding['attention_mask'].squeeze(0),  # 删除批处理维度
                'labels': torch.tensor(self.labels[idx], dtype=torch.long)}


# 初始化分词器和模型
tokenizer = BertTokenizer.from_pretrained(r"D:\codes\torch_learn_pro\bert_model")
model = BertForSequenceClassification.from_pretrained(r"D:\codes\torch_learn_pro\bert_model")

# 创建数据集
train_file = r'D:\mrpc_data\msr_paraphrase_train.txt'  # 替换为您的实际文件路径
test_file = r'D:\mrpc_data\msr_paraphrase_test.txt'  # 替换为您的实际文件路径
train_dataset = MRPCDataset(tokenizer, train_file)
test_dataset = MRPCDataset(tokenizer, test_file)

# 创建数据加载器
train_loader = DataLoader(train_dataset, batch_size=8, shuffle=True)
test_loader = DataLoader(test_dataset, batch_size=8, shuffle=False)

# 选择设备
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model.to(device)

# 创建优化器
optimizer = AdamW(model.parameters(), lr=1e-5)

# 训练模型
model.train()
for epoch in range(3):  # 多次循环以增加负载
    for batch in train_loader:
        batch = {k: v.to(device) for k, v in batch.items()}
        outputs = model(**batch)

        # 计算损失并进行反向传播
        loss = outputs.loss
        loss.backward()
        optimizer.step()
        optimizer.zero_grad()

运行结果：

注意事项：在探讨在进行压力测试时应注意的关键问题，包括内存溢出、GPU优化、数据预处理和模型的微调等，另外保证风量充足，电源供电稳定，同时开启smi进行观测记录这个测试观测。