HuggingFace Transformers语义分割实战指南-优快云博客

HuggingFace Transformers语义分割实战指南

【免费下载链接】transformers huggingface/transformers: 是一个基于 Python 的自然语言处理库，它使用了 PostgreSQL 数据库存储数据。适合用于自然语言处理任务的开发和实现，特别是对于需要使用 Python 和 PostgreSQL 数据库的场景。特点是自然语言处理库、Python、PostgreSQL 数据库。项目地址: https://gitcode.com/GitHub_Trending/tra/transformers

语义分割是计算机视觉中的一项重要任务，它要求模型对图像中的每个像素进行分类。本文将详细介绍如何使用HuggingFace Transformers库进行语义分割模型的训练和推理。

语义分割简介

语义分割（Semantic Segmentation）是计算机视觉中的一项像素级分类任务，与目标检测不同，它不仅识别图像中的物体，还需要精确到每个像素的类别划分。这项技术在自动驾驶、医学图像分析、遥感图像处理等领域有广泛应用。

准备工作

自定义数据集准备

要使用自定义数据集进行训练，需要完成以下两个关键步骤：

创建DatasetDict：需要构建包含"image"和"label"两列的数据集，两列都应为Image类型。以下是创建示例：

from datasets import Dataset, DatasetDict, Image

def create_dataset(image_paths, label_paths):
    dataset = Dataset.from_dict({
        "image": sorted(image_paths),
        "label": sorted(label_paths)
    })
    dataset = dataset.cast_column("image", Image())
    dataset = dataset.cast_column("label", Image())
    return dataset

创建id2label映射：需要准备一个JSON文件，定义类别ID到类别名称的映射关系：

import json
id2label = {0: "道路", 1: "人行道", 2: "建筑物"}
with open("id2label.json", "w") as f:
    json.dump(id2label, f)

模型训练方法

Transformers提供了两种训练语义分割模型的方式：

1. 使用Trainer API（推荐）

Trainer API封装了训练过程的复杂性，支持分布式训练、混合精度训练等高级功能。以下是使用Trainer训练SegFormer模型的示例命令：

python run_semantic_segmentation.py \
    --model_name_or_path nvidia/mit-b0 \
    --dataset_name segments/sidewalk-semantic \
    --output_dir ./output \
    --learning_rate 6e-5 \
    --per_device_train_batch_size 8 \
    --max_steps 10000 \
    --push_to_hub

关键参数说明：

model_name_or_path: 预训练模型名称或路径
dataset_name: 数据集名称
learning_rate: 学习率，SegFormer论文推荐6e-5
per_device_train_batch_size: 每个设备的训练批次大小

2. 使用Accelerate（自定义训练循环）

对于需要更灵活控制训练过程的用户，可以使用Accelerate实现自定义训练循环：

accelerate launch run_semantic_segmentation_no_trainer.py \
    --output_dir segformer-finetuned-sidewalk \
    --with_tracking \
    --push_to_hub

使用前需要先配置训练环境：

accelerate config
accelerate test

模型推理

训练完成后，可以轻松加载模型进行推理：

from transformers import AutoImageProcessor, AutoModelForSemanticSegmentation

model_name = "训练好的模型路径"
image_processor = AutoImageProcessor.from_pretrained(model_name)
model = AutoModelForSemanticSegmentation.from_pretrained(model_name)

推理示例代码：

from PIL import Image
import torch
import torch.nn as nn

image = Image.open("测试图片.jpg")
inputs = image_processor(images=image, return_tensors="pt")

with torch.no_grad():
    outputs = model(**inputs)
    
# 将输出调整到原始图像尺寸
logits = nn.functional.interpolate(
    outputs.logits.detach().cpu(),
    size=image.size[::-1],  # (高度, 宽度)
    mode="bilinear",
    align_corners=False
)

predicted = logits.argmax(1)  # 获取每个像素的预测类别

重要注意事项

背景类处理：某些数据集（如scene_parse_150）包含背景类（通常标记为0）。在训练时，应将背景类替换为255（PyTorch损失函数的ignore_index），并将其他标签减1。
标签缩减：对于包含背景类的数据集，需要设置do_reduce_labels参数为True。
训练时长：语义分割模型通常需要较长时间训练才能达到良好效果。例如，SegFormer论文中在ADE20K数据集上训练了160,000步。
学习率调度：推荐使用多项式学习率调度器（polynomial LR scheduler），这是许多语义分割模型的标配。

通过本指南，您应该能够使用HuggingFace Transformers库轻松训练和部署语义分割模型。无论是使用简单的Trainer API还是灵活的自定义训练循环，都能满足不同场景下的需求。

创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考