SAM2 多类别分割模型训练与推理指南

模型简介

本项目基于 Facebook SAM2.1-Hiera-Large 模型，针对多类别语义分割任务进行了定制化训练。模型支持 4 类分割：

类别 ID	颜色 (RGB)	描述
0	(0, 0, 127)	深蓝色 - 背景
1	(0, 212, 255)	淡蓝色
2	(127, 0, 0)	深红色
3	(255, 229, 0)	黄色

模型架构：

基础模型：transformers.Sam2Model (facebook/sam2.1-hiera-large)
视觉编码器：Hiera-Large (144→288→576 维)
自定义分类器头：Conv2d + BatchNorm + ReLU + Conv2d
输出：多类别分割 mask (256×256)

训练效果：

30 个训练周期后，平均交并比（mIoU）达到 96.51%
损失值从 0.299 降至 0.091（下降 69.5%）

运行环境要求

硬件要求

NPU 设备：华为昇腾 NPU（推荐 Ascend 910）
内存：≥ 32GB
存储：≥ 10GB（用于存储模型权重和数据集）

软件环境

Python 3.11+
PyTorch 2.7.1+（NPU 版本）
torch_npu
transformers ≥ 4.40
PIL、numpy

容器下载

docker pull swr.cn-southwest-2.myhuaweicloud.com/atelier/pytorch_ascend:pytorch_2.7.1-cann_8.3.rc1-py_3.11-hce_2.0.2509-aarch64-snt9b-20260329090059-baf3933

运行容器命令

mkdir -p /data/sam2-train-data
export IMAGE=swr.cn-southwest-2.myhuaweicloud.com/atelier/pytorch_ascend:pytorch_2.7.1-cann_8.3.rc1-py_3.11-hce_2.0.2509-aarch64-snt9b-20260329090059-baf3933
export NAME=sam2
docker run -u root --privileged \
    --name $NAME \
    --net=host \
    --shm-size=16g \
    --device /dev/davinci0 \
    --device /dev/davinci1 \
    --device /dev/davinci2 \
    --device /dev/davinci3 \
    --device /dev/davinci4 \
    --device /dev/davinci5 \
    --device /dev/davinci6 \
    --device /dev/davinci7 \
    --device /dev/davinci_manager \
    --device /dev/devmm_svm \
    --device /dev/hisi_hdc \
    -v /usr/local/dcmi:/usr/local/dcmi \
    -v /usr/local/Ascend/driver/tools/hccn_tool:/usr/local/Ascend/driver/tools/hccn_tool \
    -v /usr/local/bin/npu-smi:/usr/local/bin/npu-smi \
    -v /usr/local/Ascend/driver/lib64/:/usr/local/Ascend/driver/lib64/ \
    -v /usr/local/Ascend/driver/version.info:/usr/local/Ascend/driver/version.info \
    -v /etc/ascend_install.info:/etc/ascend_install.info \
    -v /etc/hccn.conf:/etc/hccn.conf \
    -v /data/sam2-train-data:/data/sam2-train-data \
    -it $IMAGE bash

环境设置

# 设置 Ascend 环境（容器运行可选）
source /usr/local/Ascend/ascend-toolkit/set_env.sh

# 设置线程限制（容器运行可选）
export OPENBLAS_NUM_THREADS=1

# 设置 HuggingFace 镜像 (国内)
export HF_ENDPOINT=https://hf-mirror.com

数据集格式

数据集采用分层目录结构：

data_dir/
├── JPEGImages/
│   ├── category_1/
│   │   ├── image_001.png
│   │   ├── image_002.png
│   │   └── ...
│   ├── category_2/
│   │   └── ...
│   └── ...
└── Annotations/
    ├── category_1/
    │   ├── image_001.png  # 彩色标注
    │   ├── image_002.png
    │   └── ...
    ├── category_2/
    │   └── ...
    └── ...

标注图片要求：

格式：PNG (RGB)
颜色必须与固定映射一致：
- 背景：RGB(0, 0, 127)
- 前景1：RGB(0, 212, 255)
- 前景2：RGB(127, 0, 0)
- 前景3：RGB(255, 229, 0)

训练命令

单卡训练

从头开始训练

python run_train_npu_multiclass.py \
    --model facebook/sam2.1-hiera-large \
    --data_dir /data/sam2-train-data \
    --output_dir ./output_fixed \
    --epochs 30 \
    --batch_size 1 \
    --base_lr 1e-5 \
    --vision_lr 1e-5 \
    --device npu:0

从 checkpoint 续训

python run_train_npu_multiclass.py \
    --checkpoint ./output_fixed/best_model.pt \
    --model facebook/sam2.1-hiera-large \
    --data_dir /data/sam2-train-data \
    --output_dir ./output_fixed \
    --epochs 60 \
    --batch_size 1 \
    --base_lr 1e-5 \
    --vision_lr 1e-5 \
    --device npu:0

后台训练

nohup python run_train_npu_multiclass.py \
    --data_dir /data/sam2-train-data \
    --output_dir ./output_fixed \
    --epochs 30 \
    > train.log 2>&1 &

# 查看训练日志
tail -f train.log

多卡分布式训练

使用 torchrun 启动 (推荐)

# 8 卡分布式训练
torchrun --nproc_per_node=8 \
    run_train_npu_multiclass.py \
    --model facebook/sam2.1-hiera-large \
    --data_dir /data/sam2-train-data \
    --output_dir ./output_fixed \
    --epochs 30 \
    --batch_size 1 \
    --base_lr 1e-5 \
    --vision_lr 1e-5 \
    --num_workers 0

使用 torchrun 续训

torchrun --nproc_per_node=8 \
    run_train_npu_multiclass.py \
    --checkpoint ./output_fixed/best_model.pt \
    --model facebook/sam2.1-hiera-large \
    --data_dir /data/sam2-train-data \
    --output_dir ./output_fixed \
    --epochs 60 \
    --batch_size 1 \
    --base_lr 1e-5 \
    --vision_lr 1e-5 \
    --num_workers 0

使用分布式启动脚本

# 创建启动脚本
cat > distributed_train.sh << 'EOF'
#!/bin/bash
source /usr/local/Ascend/ascend-toolkit/set_env.sh
export OPENBLAS_NUM_THREADS=1
export HF_ENDPOINT=https://hf-mirror.com

torchrun --nproc_per_node=8 \
    run_train_npu_multiclass.py \
    --model facebook/sam2.1-hiera-large \
    --data_dir /data/sam2-train-data \
    --output_dir ./output_fixed \
    --epochs 30 \
    --batch_size 1 \
    --base_lr 1e-5 \
    --vision_lr 1e-5 \
    --num_workers 0
EOF

chmod +x distributed_train.sh

# 后台运行分布式训练
nohup ./distributed_train.sh > distributed_train.log 2>&1 &

指定 GPU/NPU 数量

# 使用 4 卡
torchrun --nproc_per_node=4 run_train_npu_multiclass.py ...

# 使用 2 卡
torchrun --nproc_per_node=2 run_train_npu_multiclass.py ...

分布式训练注意事项

batch_size 会自动乘以 GPU/NPU 数量：实际 batch_size = batch_size × nproc_per_node
只在 rank 0 保存 checkpoint 和打印日志
使用 DistributedSampler 自动分配数据
学习率建议相应调大：分布式训练时可将 base_lr 调大为单卡的 n 倍
建议设置 --num_workers 0：避免共享内存不足导致训练中断
环境变量自动设置：torchrun 会自动设置 LOCAL_RANK、RANK、WORLD_SIZE
HCCL 仅支持 Float32：分布式通信时必须使用 Float32，不支持 Float64

参数说明

参数	默认值	说明
`--model`	facebook/sam2.1-hiera-large	预训练模型 ID
`--data_dir`	必填	数据集目录
`--output_dir`	./output	输出目录
`--epochs`	10	训练轮数
`--batch_size`	1	单卡批大小
`--base_lr`	1e-4	基础学习率 (分类器等)
`--vision_lr`	1e-5	Vision Encoder 学习率
`--checkpoint`	None	续训 checkpoint 路径
`--device`	npu:0	设备 (单卡训练时)
`--local_rank`	0	分布式训练自动设置
`--num_workers`	4	DataLoader workers 数量，分布式建议设为 0

训练指标分析

使用分析脚本

python analyze_training.py --output_dir ./output_fixed/

输出示例

Training Progress Summary:
============================================================
Epoch      Loss            mIoU            Status
------------------------------------------------------------
5          0.2990          0.9037          Saved
10         0.1486          0.9593          Saved
15         0.1363          0.9461          Saved
20         0.1054          0.9590          Saved
25         0.0941          0.9637          Saved
30         0.0912          0.9651          Saved
============================================================

Best Model:
------------------------------------------------------------
  File:     best_model.pt
  Epoch:    30
  Loss:     0.0912
  mIoU:     0.9651
------------------------------------------------------------

Training Analysis:
------------------------------------------------------------
Loss improvement: 69.5% (0.2990 -> 0.0912)
mIoU improvement: 6.8% (0.9037 -> 0.9651)
Best model: Epoch 30, mIoU = 0.9651, Loss = 0.0912

指标说明

指标	说明	计算方式
Loss	总损失	CrossEntropy + 0.5×Dice Loss
mIoU	平均交并比	所有类别 IoU 的平均值
Epoch Time	单轮训练时间	NPU 推理 + 数据加载时间
LR	当前学习率	Warmup + Cosine Annealing

正常训练曲线特征

Loss：持续下降，后期趋于平稳
mIoU：持续上升，后期趋于平稳
LR：Warmup (5 epochs) 线性增长，之后 Cosine Annealing 下降

推理命令

基本推理

python run_inference_final.py \
    --checkpoint ./output_fixed/best_model.pt \
    --image /path/to/image.png \
    --output mask.png \
    --device npu:0

指定分割点

python run_inference_final.py \
    --checkpoint ./output_fixed/best_model.pt \
    --image /path/to/image.png \
    --output mask.png \
    --point 600,600

使用边界框分割

python run_inference_final.py \
    --checkpoint ./output_fixed/best_model.pt \
    --image /path/to/image.png \
    --output mask.png \
    --box 100,100,500,500

输出灰度 mask

python run_inference_final.py \
    --checkpoint ./output_fixed/best_model.pt \
    --image /path/to/image.png \
    --output mask.png \
    --no_color

推理参数说明

参数	说明
`--checkpoint`	训练好的模型路径
`--image`	输入图片路径
`--output`	输出 mask 路径
`--device`	设备 (npu:0 或 cpu)
`--point`	分割点坐标 (x,y)，默认图片中心
`--box`	边界框 (x_min,y_min,x_max,y_max)
`--no_color`	输出灰度 mask
`--model`	预训练模型 ID

推理输出示例

============================================================
SAM2 NPU Inference - Multi-class (Fixed Version)
============================================================
Loading checkpoint from: ./output_fixed/best_model.pt
Checkpoint epoch: 30
Checkpoint IoU: 0.9651
Number of classes: 4

各类别占比:
  深蓝色(背景): 20.7%
  淡蓝色: 32.3%
  深红色: 26.4%
  黄色: 20.6%

Mask saved to: mask.png
============================================================

输出文件说明

训练输出目录

output_fixed/
├── best_model.pt           # 最佳模型 (mIoU 最高)
├── checkpoint_epoch_5.pt   # 第 5 轮 checkpoint
├── checkpoint_epoch_10.pt  # 第 10 轮 checkpoint
├── checkpoint_epoch_15.pt  # 第 15 轮 checkpoint
├── checkpoint_epoch_20.pt  # 第 20 轮 checkpoint
├── checkpoint_epoch_25.pt  # 第 25 轮 checkpoint
├── checkpoint_epoch_30.pt  # 第 30 轮 checkpoint
└── ...

检查点（Checkpoint）内容

每个检查点包含：

epoch: 训练轮数
model_state_dict: 模型权重
optimizer_state_dict: 优化器状态
loss: 平均损失
iou: 平均交并比（mIoU）
num_classes: 类别数量 (4)
num_frames: 视频帧数 (1)
max_num_objects: 最大对象数 (10)
base_lr: 基础学习率
vision_lr: 视觉编码器（Vision Encoder）学习率
total_epochs: 总训练轮数
color_to_class: 颜色映射信息

常见问题

1. 续训后性能下降

原因：优化器（Optimizer）状态未正确恢复或模型架构不匹配
解决：确保使用 --model facebook/sam2.1-hiera-large 参数

2. 推理颜色与标注不一致

原因：颜色映射未固定
解决：使用修复后的训练脚本重新训练

3. NPU 内存不足

解决：减小批次大小（batch_size）或使用梯度检查点（gradient checkpointing）

4. HuggingFace 模型下载失败

解决：设置镜像 export HF_ENDPOINT=https://hf-mirror.com

5. 分布式训练初始化失败

原因：环境变量未正确设置
解决：确保已执行 source /usr/local/Ascend/ascend-toolkit/set_env.sh

6. 分布式训练时检查点（checkpoint）加载错误

原因：DDP 包装后模型结构与保存时不同
解决：脚本已自动处理，使用 model.module.state_dict() 保存

7. HCCL 聚合（allreduce）报错：Unsupported data type at::kDouble

原因：HCCL 通信库不支持 Float64（双精度）类型，仅支持 Float32/Float16
解决：已修复，分布式训练时使用 dtype=torch.float32 创建张量（tensor）

错误示例：

RuntimeError: HCCL allreduce: Unsupported data type at::kDouble
ERR02007 DIST feature not supported

修复位置：run_train_npu_multiclass.py 第 350-352 行

# 修复前（错误）
loss_tensor = torch.tensor([avg_loss], device=device)  # 默认 Float64

# 修复后（正确）
loss_tensor = torch.tensor([avg_loss], dtype=torch.float32, device=device)
iou_tensor = torch.tensor([avg_iou], dtype=torch.float32, device=device)

文件清单

SAM2-A2/
├── run_train_npu_multiclass.py  # 训练脚本 (支持分布式，已修复HCCL问题)
├── run_inference_final.py       # 推理脚本
├── analyze_training.py          # 训练分析脚本
└── output_fixed/                # 输出目录
    ├── best_model.pt
    └── checkpoint_epoch_*.pt

版本历史

版本	日期	更新内容
v1.0	2026-04-21	初始版本，固定颜色映射，4 类分割
v1.1	2026-04-21	修复续训 optimizer 恢复问题
v1.2	2026-04-21	添加分布式训练支持 (torchrun)
v1.3	2026-05-14	修复 HCCL Float64 类型不支持问题

联系方式

如有问题，请查看训练日志或使用 analyze_training.py 分析训练状态。