基于BERT的智能商品分类系统

项目概述

在电商平台中，商品发布通常是一个繁琐的过程，需要手动录入商品的各项信息，如品牌、品类等。为了提升效率并减轻商家的工作负担，智能商品录入系统应运而生。本项目是一个基于BERT的智能商品分类系统，包含数据预处理、模型训练、评估、预测和Web服务五个核心流程。分别运行命令python main.py 参数运行不同的模块。

if __name__ == '__main__':
    parser = ArgumentParser()
    parser.add_argument("action",choices=['train','predict','evaluate','preprocess','serve'])
    args = parser.parse_args()
    action = args.action
    if action == "preprocess":
        from preprocess.preprocess import preprocess
        preprocess()
    elif action == "train":
        from runner.train import train
        train()
    elif action == 'predict':
        from runner.predict import predict
        predict()
    elif action == 'evaluate':
        from runner.evaluate import evaluate
        evaluate()
    elif action == 'serve':
        from web.app import serve
        serve()

技术栈

➢ pytorch：深度学习框架，用于训练和推理
➢ transformers：Hugging Face 提供的库，用于加载和微调 BERT 等预训练模型。
➢ datasets：用于高效加载和处理大规模数据集。
➢ scikit-learn：用于模型评估。
➢ tensorboard：用于可视化训练过程中的损失、准确率等指标。
➢ tqdm：用于显示训练进度条，方便监控训练过程。
➢ jupyter：用于实验和数据分析。
➢ FastAPI：用于构建和部署 API 接口。
➢ Uvicorn：FastAPI 的服务器，用于高性能地运行 FastAPI 应用。

数据预处理

数据集说明

本项目使用的 商品标题分类数据集 来自 百度 AI Studio，该数据集包含了训练集、验证集和测试集。数据结构如下：
在这里插入图片描述

在这里插入图片描述

1.数据集提供了train、test、valid三个分开的文件。
2.训练集：88824条，验证集：22207条，测试集：22207条
3.数据集包含两个字段：label和texta，分别表示分类标签和标题数据。

预处理过程

graph TD A[preprocess.py] --> B[加载原始数据] B -->|读取 train.txt/valid.txt/test.txt| C[dataset_dict] C --> D[label编码处理] D -->|获取所有分类名称| E[ClassLabel转换] E --> F[加载BERT分词器] F --> G[batch_encode函数] G -->|tokenizer分词| H[生成input_ids/attention_mask] H --> I[添加labels字段] I --> J[移除原始label/text_a列] J --> K[保存到processed目录] style A fill:#f9f,stroke:#333,stroke-width:2px

def preprocess():
    # 1.从文件中加载数据集
    dataset_dict = load_dataset(
        'csv',
        data_files={
            'train': str(RAW_DATA_DIR/RAW_TRAIN_DATA),
            'valid': str(RAW_DATA_DIR/RAW_VALID_DATA),
            'test': str(RAW_DATA_DIR/RAW_TEST_DATA)
        },
        delimiter='\t'
    )
    print(dataset_dict)
    # 2.label编码处理
    # 2.1获取所有分类的名称，set去重
    all_labels = sorted(set(dataset_dict['train']['label']))
    # 2.2列转换编码
    dataset_dict = dataset_dict.cast_column('label',ClassLabel(names=all_labels))
    print(dataset_dict['train'][0:2])
    '''输出结果
    {'label': [9, 10], 'text_a': ['樱之歌蓝色之恋5件套日式釉下彩纯手绘家用餐具套装陶瓷器碗盘碗碟微波炉可用', '北欧欧慕料理机手持料理棒家用多功能手持式搅拌机打蛋切菜婴儿辅食绞肉机榨汁机NJB103SA冰湖蓝']}
    '''
    # 3.加载器分词器
    tokenizer = AutoTokenizer.from_pretrained(BERT_MODEL_NAME)
    # 4.处理标题文本，得到模型输入
    def batch_encode(examples):
        inputs = tokenizer(examples['text_a'],truncation=True)
        inputs['labels'] = examples['label']
        return inputs
    dataset_dict = dataset_dict.map(batch_encode,batched=True,remove_columns=['label','text_a'])
    print(dataset_dict['train'][0:2])
    # 5.保存数据集
    dataset_dict.save_to_disk(PROCESSED_DATA_DIR)

程序首先从data/raw/目录下读取三个原始数据文件：train.txt、valid.txt、test.txt，这些文件是制表符分隔的CSV格式，包含text_a（商品标题）和label（分类标签）两列；
使用Hugging Face的datasets库将数据加载为DatasetDict对象，包含train、valid、test三个数据集；
获取训练集中所有label并去重排序，生成类别列表all_labels；
将label列转换为ClassLabel类型，实现字符串标签到整数ID的映射；
加载预训练的BERT中文分词器（bert-base-chinese）；
定义batch_encode函数，对每条数据的text_a字段进行分词，生成模型所需的input_ids和attention_mask，同时将label转换为labels字段；
使用map方法批量处理整个数据集，并移除原始的label和text_a列；

模型训练

主要流程

flowchart TD A["torch.device() 检测CUDA设备"] B["AutoTokenizer.from_pretrained() 加载BERT分词器"] C["读取LABELS_FILE 构建id2label/label2id"] D["AutoModelForSequenceClassification.from_pretrained() 加载预训练BERT分类模型"] E["get_dataset('train') 加载训练数据集"] F["get_dataset('valid') 加载验证数据集"] G["DataCollatorWithPadding() 动态填充数据整理器"] H["compute_metrics() 定义准确率与F1评估指标"] I["TrainConfig() 配置训练参数(batch_size=8)"] J["Trainer() 初始化训练器"] K["trainer.train() 执行训练循环"] A --> B --> C --> D --> E --> F --> G --> H --> I --> J --> K

def train():
    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
    tokenizer = AutoTokenizer.from_pretrained(BERT_MODEL_NAME)
    with open(MODEL_DIR/LABELS_FILE,'r',encoding='utf-8') as f:
        all_labels = f.read().splitlines()
    id2label = {index:label for index,label in enumerate(all_labels)}
    label2id = {label:index for index,label in enumerate(all_labels)}
    model = AutoModelForSequenceClassification.from_pretrained(
        BERT_MODEL_NAME,
        num_labels=len(all_labels),
        id2label=id2label,
        label2id=label2id
    )
    print(model.config.id2label)
    # model.save_pretrained(MODEL_DIR)
    train_dataset = get_dataset('train')
    valid_dataset = get_dataset('valid')
    collate_fn = DataCollatorWithPadding(
        tokenizer=tokenizer,
        padding=True,
        return_tensors='pt'
    )
    def compute_metrics(preds,labels):
        acc = accuracy_score(labels,preds)
        f1 = f1_score(labels,preds,average='weighted')
        return {'acc': acc,'f1': f1}
    train_config = TrainConfig(batch_size=8,output_dir=MODEL_DIR,log_dir=LOG_DIR)
    trainer = Trainer(
        model=model,
        train_dataset=train_dataset,
        valid_dataset=valid_dataset,
        collate_fn=collate_fn,
        compute_metrics=compute_metrics,
        device=device,
        train_config=train_config
    )
    trainer.train()

def get_dataset(ds_type="train"):
    path = str(PROCESSED_DATA_DIR/ds_type)
    dataset = load_from_disk(path)
    return dataset

开始训练

该训练流程首先加载历史检查点恢复训练状态，并将模型设置为训练模式，随后获取训练数据加载器。在每个训练轮次中，逐批次执行单步训练并累计步数，每经过save_steps步会记录当前损失、执行模型评估并输出评估指标，同时根据早停策略判断是否提前终止训练，若满足停止条件则直接返回，否则保存当前检查点后继续迭代，直至所有轮次完成。

flowchart TD A[开始训练] --> B[加载检查点 _load_checkpoint] B --> C[设置模型为训练模式 model.train] C --> D[获取数据加载器 _get_dataloader] D --> E[遍历 epoch 循环] E --> F[遍历 dataloader 批次数据] F --> G[执行单步训练 _train_one_step] G --> H{step % save_steps == 0?} H -->|否| I[step += 1] I --> F H -->|是| J[记录损失到 TensorBoard] J --> K[执行评估 evaluate] K --> L{_should_stop 判断是否早停?} L -->|是| M[输出早停信息并返回] L -->|否| N[保存检查点 _save_checkpoint] N --> I M --> O[结束训练] E -->|所有epoch完成| O

class Trainer:
    ...
    def train(self):
        # 加载检查点
        self._load_checkpoint()
        self.model.train()
        dataloader = self._get_dataloader(self.train_dataset)
        for epoch in range(self.train_config.epochs):
            for inputs in tqdm(dataloader,desc=f'[Epoch:{epoch+1}]'):
                this_loss = self._train_one_step(inputs)
                if self.step % self.train_config.save_steps == 0:
                    tqdm.write(f'[Epoch:{epoch+1}|Step:{self.step}] Loss:{this_loss}')
                    self.writer.add_scalar('loss', this_loss, self.step)
                    metrics = self.evaluate()
                    metrics_str = '|'.join([f'{k}:{v:.4f}' for k, v in metrics.items()])
                    tqdm.write(f'[Evaluate:{metrics_str}]')
                    if self._should_stop(metrics):
                        tqdm.write('早停')
                        return
                    self._save_checkpoint()
                self.step += 1

获取数据加载器

将数据集转换为 PyTorch 中可以使用的 Dataset 和 DataLoader，可以使用torch.utils.data.DataLoader来加载数据集并进行批量处理。DataLoader需要一个整理函数来处理批量数据，例如填充（padding）。collate_fn使用transformers库中的data_collator.with_padding函数，该函数接受tokenizer作为参数，并执行填充操作。

class Trainer:
    ...
    def _get_dataloader(self,dataset):
        dataset.set_format(type="torch")
        dataloader = DataLoader(
            dataset,
            batch_size=self.train_config.batch_size,
            shuffle = True,
            collate_fn = self.collate_fn
        )
        return dataloader

执行单步训练

首先将输入数据张量转移到指定计算设备上，随后在autocast混合精度上下文中执行模型前向传播获得损失值；接着利用scaler.scale对损失进行缩放并调用backward完成反向传播，以此防止梯度下溢；之后通过scaler.step先反缩放梯度并检查是否存在inf或NaN，若无异常则执行optimizer.step更新模型参数，否则跳过本次更新；最后调用scaler.update动态调整缩放因子，并清空梯度以准备下一步迭代，最终返回当前批次的损失值。

class Trainer:
    ...
    def _train_one_step(self,inputs):
        inputs = {k: v.to(self.device) for k, v in inputs.items()}
        # 使用 autocast 进行前向传播，实现混合精度加速
        with torch.autocast(
                device_type=self.device.type,
                dtype=torch.float16,
                enabled=self.train_config.use_amp,
        ):
            outputs = self.model(**inputs)
            loss = outputs.loss
        # 缩放 loss，调用 backward() 时对 loss 进行缩放，生成缩放后的梯度，防止梯度下溢
        self.scaler.scale(loss).backward()
        # scaler.step() 会先对优化器中参数的梯度进行反缩放（恢复正常范围）
        # 如果梯度中不包含 inf 或 NaN，才会执行 optimizer.step() 更新参数
        # 如果包含 inf 或 NaN， optimizer.step() 会被跳过，避免错误更新
        self.scaler.step(self.optimizer)
        # 更新 scaler 的缩放因子，供下一次迭代使用，自动动态调整 scale 保证稳定性
        self.scaler.update()
        self.optimizer.zero_grad()
        return loss.item()

混合精度训练

混合精度训练是指在深度学习模型训练时，一些操作使用 torch.float32（单精度）数据
类型以保证数值稳定性，而其他操作使用 torch.float16（半精度）以减少显存占用和加快计
算速度，从而在不影响模型精度和收敛效果的前提下，实现更高效的训练。
在 PyTorch 中，混合精度训练通常使用自动混合精度（Automatic Mixed Precision, AMP）
模块实现，主要依赖以下两个组件：
➢ torch.autocast
➢ torch.cuda.amp.GradScaler

1）torch.autocast
torch.autocast 是 PyTorch 自动混合精度（AMP）的上下文管理器，用于自动选择适当
的精度进行计算。例如，在卷积、矩阵乘法等可使用半精度加速的操作中使用 torch.float16，
而在归一化、损失计算等需要高精度的操作中保持 torch.float32，从而在节省显存和加速训
练的同时保证数值稳定性。

2）torch.cuda.amp.GradScaler
torch.cuda.amp.GradScaler 是用于动态损失缩放（Dynamic Loss Scaling）的工具，用来
防止在使用半精度训练时因梯度值过小导致的下溢（underflow），保证训练过程稳定收敛。
它会动态调整缩放因子，将损失放大后再进行反向传播以避免梯度消失，在优化器更新时再
自动缩小回原始范围，同时监控是否出现溢出并动态调整缩放值。

执行评估

1.训练过程中只保存最小损失的模型可能导致过拟合。
2.保存模型的标准应改为验证集上的表现。
3.早停法：当验证损失不再下降时，停止训练并保存当前最佳模型。
4.早停法的耐心参数patience控制容忍度，连续多次损失不下降时停止训练。
5.调用Trainer的evaluate方法可以遍历验证集数据，获取不同的评价指标。

class Trainer:
    ...
    def evaluate(self):
        dataloader = self._get_dataloader(self.valid_dataset)
        self.model.eval()
        total_loss = 0.0
        all_labels = []
        all_preds = []
        for inputs in tqdm(dataloader,desc=f'[Evaluate]'):
            inputs = {k: v.to(self.device) for k,v in inputs.items()}
            outputs = self.model(**inputs)
            loss = outputs.loss
            total_loss += loss.item()
            logits = outputs.logits
            preds = torch.argmax(logits,dim=-1)
            all_preds.extend(preds.tolist())
            labels = inputs['labels']
            all_labels.extend(labels.tolist())
        loss = total_loss / len(dataloader)
        metrics = self.compute_metrics(all_preds,all_labels)
        return {'loss': loss,**metrics}

判断早停

class Trainer:
    ...
    def _should_stop(self,metrics):
        metric = metrics[self.train_config.early_stop_metric]
        score = -metric if self.train_config.early_stop_metric == 'loss' else metric
        if score > self.early_stop_best_score:
            self.early_stop_best_score = score
            self.early_stop_counter = 0
            tqdm.write('保存最佳模型...')
            self.model.save_pretrained(str(Path(self.train_config.output_dir)/"best"))
            return False
        else:
            self.early_stop_counter += 1
            if self.early_stop_counter >= self.train_config.early_stop_patience:
                return True
            else:
                return False

在这里插入图片描述

模型预测

首先初始化计算设备并加载分词器，随后加载训练好的最优分类模型创建预测器对象。调用predict方法时，先将输入统一转换为列表格式，经过分词器编码为张量后送入模型进行前向推理，从输出的logits中取最大值的类别ID，再通过id2label映射解码为对应的中文标签，最终根据原始输入类型返回单个标签或标签列表。

graph TD A[predict.py] --> B[初始化设备和Tokenizer] B --> C[加载best模型] C --> D[创建Predictor] D --> E[调用predict进行预测] E --> F[输入统一转为列表] F --> G[Tokenizer编码为张量] G --> H[模型前向传播获取logits] H --> I[argmax获取类别ID] I --> J[解码为中文标签] J --> K[根据输入类型返回单个或列表]

class Predictor:
    # 初始化，传入model，tokenizer，device
    def __init__(self,model,tokenizer,device):
        self.model = model
        self.tokenizer = tokenizer
        self.device = device
        self.model.to(device)
    # 核心预测方法
    def predict(self,texts):
        is_str = isinstance(texts,str)
        # 统一数据格式：如果是字符串，转换为列表
        if is_str:
            texts = [texts]
        # 1.分词编码，得到模型输入
        inputs = self.tokenizer(texts,padding=True,truncation=True,return_tensors='pt')
        inputs = {k: v.to(self.device) for k,v in inputs.items()}
        # 2.前向传播
        with torch.no_grad():
            outputs = self.model(**inputs)
        # 3.根据输出结果，解码得到中文分类标签
        preds = torch.argmax(outputs.logits,dim=-1).tolist()
        labels = [self.model.config.id2label[pred_id] for pred_id in preds]
        if is_str:
            return labels[0]
        return labels
def predict():
    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
    tokenizer = AutoTokenizer.from_pretrained(BERT_MODEL_NAME)
    # 加载微调后的模型
    model = AutoModelForSequenceClassification.from_pretrained(MODEL_DIR/'best')
    predictor = Predictor(model,tokenizer,device)
    text = '好奇心钻装纸尿裤L40片9-14kg'
    result = predictor.predict(text)
    print(result)
    texts = ["240ML*15养元2430六个核桃","潘婷丝质顺滑洗发露750ml","640G正航牛奶早餐饼干"]
    result = predictor.predict(texts)
    print(result)

在这里插入图片描述

模型评估

1.模型评估的目的是评估模型的性能，确定最优模型。
2.通过调整容忍度，控制模型的训练过程，防止过早停止。
3.当容忍度设置为较大值时，模型会继续训练，直到达到最优或早停条件。
4.评估器（Evaluator）用于对训练好的模型进行评估。
5.评估器可以调用训练器（Trainer）的evaluate方法进行评估。
6.评估器接收模型、数据集和配置等参数，并返回评估结果。

# 验证流程
def evaluate():
    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
    tokenizer = AutoTokenizer.from_pretrained(BERT_MODEL_NAME)
    model = AutoModelForSequenceClassification.from_pretrained(MODEL_DIR/'best')
    test_dataset = get_dataset('test')
    collate_fn = DataCollatorWithPadding(
        tokenizer=tokenizer,
        padding=True,
        return_tensors='pt'
    )
    def compute_metrics(preds,labels):
        acc = accuracy_score(labels,preds)
        f1 = f1_score(labels,preds,average='weighted')
        return {'acc': acc,'f1': f1}
    train_config = TrainConfig(batch_size=8, output_dir=MODEL_DIR, log_dir=LOG_DIR)
    trainer = Trainer(
        model=model,
        valid_dataset=test_dataset,
        collate_fn=collate_fn,
        compute_metrics=compute_metrics,
        device=device,
        train_config=train_config
    )
    metrics = trainer.evaluate()
    print(metrics)

在这里插入图片描述

Web服务

从FastAPI应用初始化开始，依次挂载静态目录、初始化计算设备、加载分词器和最优分类模型，并创建Predictor和TitleService服务对象。应用提供两个路由：GET根路径返回前端页面，POST /predict接口接收包含text字段的Title请求体，通过TitleService调用Predictor进行模型推理，最终将分类结果封装为Category响应返回给客户端。

flowchart TD A["FastAPI应用初始化 (app.py)"] --> B["挂载静态文件目录 (static)"] B --> C["初始化计算设备 (device)"] C --> D["加载分词器 (AutoTokenizer)"] D --> E["加载最佳模型 (AutoModelForSequenceClassification)"] E --> F["创建预测器 (Predictor)"] F --> G["创建标题服务 (TitleService)"] G --> H["GET / 路由"] H --> I["返回 index.html 页面"] G --> J["POST /predict 路由"] J --> K["接收 Title 请求体 (text字段)"] K --> L["调用 TitleService.predict"] L --> M["Predictor 执行推理"] M --> N["获取分类结果"] N --> O["返回 Category 响应"]

# src/web/service.py
class TitleService:
    def __init__(self,predictor):
        self.predictor = predictor
    def predict(self,title):
        return self.predictor.predict(title)
# src/web/schemas.py
from pydantic import BaseModel
class Title(BaseModel):
    text:str
class Category(BaseModel):
    category:str
# src/web/app.py
STATIC_DIR = Path(__file__).parent / "static"
app = FastAPI(title="智能商品分类")
app.mount("/static", StaticFiles(directory=STATIC_DIR), name="static")
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
tokenizer = AutoTokenizer.from_pretrained(BERT_MODEL_NAME)
model = AutoModelForSequenceClassification.from_pretrained(MODEL_DIR/'best')
# 创建预测器
predictor = Predictor(model=model,tokenizer=tokenizer,device=device)
service = TitleService(predictor=predictor)
@app.get("/")
def index():
    return FileResponse(STATIC_DIR / "index.html")
@app.post("/predict")
def predict(title:Title)->Category:
    label = service.predict(title.text)
    response = Category(category=label)
    return response
def serve():
    uvicorn.run("web.app:app",host="0.0.0.0",port=8000)

在这里插入图片描述

posted @ 2026-08-17 18:48 _丑小鸭阅读(1) 评论(0) 收藏举报

刷新页面返回顶部

xym_love

基于BERT的智能商品分类系统

项目概述

技术栈

数据预处理

数据集说明

预处理过程

模型训练

主要流程

开始训练

获取数据加载器

执行单步训练

执行评估

判断早停

模型预测

模型评估

Web服务

公告