python pytorch tensorflow transforms 模型培训脚本

环境准备

https://www.doubao.com/thread/w5e26d6401c003bb2

执行培训脚本

bash 复制代码
import torch
from torch.utils.data import Dataset, DataLoader
from transformers import DistilBertTokenizer, DistilBertForSequenceClassification, AdamW
import numpy as np


# 自定义数据集类
class SentimentDataset(Dataset):
    def __init__(self, texts, labels, tokenizer, max_length):
        self.texts = texts
        self.labels = labels
        self.tokenizer = tokenizer
        self.max_length = max_length

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

    def __getitem__(self, idx):
        text = self.texts[idx]
        label = self.labels[idx]
        encoding = self.tokenizer.encode_plus(
            text,
            add_special_tokens=True,
            max_length=self.max_length,
            padding='max_length',
            truncation=True,
            return_tensors='pt'
        )
        return {
            'input_ids': encoding['input_ids'].flatten(),
            'attention_mask': encoding['attention_mask'].flatten(),
            'labels': torch.tensor(label, dtype=torch.long)
        }


# 新增问答数据
qa_texts = ["李白是那个朝代的诗人?", "地球的卫星是什么?", "中国的首都是哪里?", "水的化学式是什么?", "苹果公司的创始人是谁?"]
qa_labels = [1, 2, 3, 4, 5]  # 为每个问题分配一个唯一的标签

# 合并数据
all_texts = qa_texts
all_labels = qa_labels

# 初始化 tokenizer
tokenizer = DistilBertTokenizer.from_pretrained('distilbert-base-uncased')

# 创建数据集和数据加载器
dataset = SentimentDataset(all_texts, all_labels, tokenizer, max_length=128)
dataloader = DataLoader(dataset, batch_size=2, shuffle=True)

# 初始化模型,注意 num_labels 需要根据总标签数调整
model = DistilBertForSequenceClassification.from_pretrained('distilbert-base-uncased', num_labels=8)

# 定义优化器,加入 weight_decay 进行正则化
optimizer = AdamW(model.parameters(), lr=1e-4, weight_decay=0.01)

# 训练模型
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
model.to(device)

# 早停策略相关参数
best_loss = float('inf')
patience = 3
counter = 0

for epoch in range(20):  # 增加训练轮数,但结合早停策略
    total_loss = 0
    model.train()
    for batch in dataloader:
        input_ids = batch['input_ids'].to(device)
        attention_mask = batch['attention_mask'].to(device)
        labels = batch['labels'].to(device)

        optimizer.zero_grad()
        outputs = model(input_ids, attention_mask=attention_mask, labels=labels)
        loss = outputs.loss
        loss.backward()
        optimizer.step()
        total_loss += loss.item()

    avg_loss = total_loss / len(dataloader)
    print(f'Epoch {epoch + 1}, Loss: {avg_loss}')

    # 早停策略
    if avg_loss < best_loss:
        best_loss = avg_loss
        counter = 0
    else:
        counter += 1
        if counter >= patience:
            print("Early stopping triggered!")
            break

# 处理提问的函数
def handle_query(model, tokenizer, query_text, device):
    model.eval()
    with torch.no_grad():
        encoding = tokenizer(query_text, return_tensors='pt', padding=True, truncation=True, max_length=128).to(device)
        logits = model(**encoding).logits
        predictions = np.argmax(logits.cpu().numpy(), axis=1)
        return predictions

# 情感标签映射
label_mapping = {
    1: "唐朝",
    2: "月球",
    3: "北京",
    4: "H2O",
    5: "史蒂夫·乔布斯、史蒂夫·沃兹尼亚克和罗恩·韦恩"
}

# 针对新增的5条数据提问
query_text = ["李白是那个朝代的诗人?", "地球的卫星是什么?", "中国的首都是哪里?", "水的化学式是什么?", "苹果公司的创始人是谁?"]
result = handle_query(model, tokenizer, query_text, device)
readable_result = [label_mapping[pred] for pred in result]
print("Query result:", readable_result)
相关推荐
海哥编程5 分钟前
Python 数据分析(二):Matplotlib 绘图
python·数据分析·matplotlib
go54631584651 小时前
Python点阵字生成与优化:从基础实现到高级渲染技术
开发语言·人工智能·python·深度学习·分类·数据挖掘
猫头虎1 小时前
2025年02月11日 Go生态洞察:Go 1.24 发布亮点全面剖析
开发语言·后端·python·golang·go·beego·go1.19
仰望天空—永强1 小时前
PS 2025【七月最新v26.5】PS铺软件安装|最新版|附带安装文件|详细安装说明|附PS插件
开发语言·图像处理·python·图形渲染·photoshop
MediaTea1 小时前
Python 库手册:xmlrpc.client 与 xmlrpc.server 模块
开发语言·python
悦悦子a啊1 小时前
Python之--字典
开发语言·python·学习
水军总督2 小时前
OpenCV+Python
python·opencv·计算机视觉
qyhua2 小时前
Windows 平台源码部署 Dify教程(不依赖 Docker)
人工智能·windows·python
一车小面包2 小时前
Python高级入门Day6
开发语言·python
攻城狮凌霄3 小时前
PHP与ChatGPT结合的技术王炸,开发高效创作小红书内容系统
python