SAM-2.1 自定义分割微调实战:从自定义数据集到自动掩码生成
这篇教程是我根据 SAM-2.1 自定义数据集微调和自动 mask 生成流程整理出来的。重点演示如何下载 SAM2 格式数据集,配置 SAM-2.1 训练脚本,训练后对比微调模型和基础模型的分割效果。
SAM-2.1 的训练流程依赖官方仓库、配置文件和 checkpoint。这个 notebook 适合作为完整项目复现模板,尤其适合希望把 SAM-2.1 迁移到特定零件、工业目标或业务数据集的场景。
本文会重点跑通以下流程:
- 下载 SAM2 格式数据集
- 克隆并安装 SAM2 官方仓库
- 修正数据集文件名以适配训练配置
- 启动 SAM-2.1 单 GPU 微调
- 对比微调模型和基础模型的自动 mask 结果
如果你正在系统学习多模态微调、目标检测、OCR 或图像分割,建议收藏本文;配套 notebook、示例图片和运行环境说明后续会继续整理。如果环境配置卡住,可以在评论区说明具体报错。
📚 文章目录
- SAM-2.1 自定义分割微调实战:从自定义数据集到自动掩码生成
- 📥 下载 SAM-2.1 与数据集
- 🧹 修改数据集文件名
- 🏋️ 开始训练 SAM-2.1
- 🔍 可视化微调模型结果
- 📌 小结
- 📚 同系列教程汇总
📥 下载 SAM-2.1 与数据集
先从数据集后台获取 car parts 数据集,再克隆 SAM2 官方仓库和训练配置。注意:原 notebook 中直接写了 API key,发布前建议替换为自己的密钥或 Colab Secrets。
fromtypesimportSimpleNamespace# 从数据集后台下载并解压数据集后,修改 DATASET_DIR 指向数据集目录。DATASET_DIR="/content/dataset"# 修改为数据集后台导出的数据集目录dataset=SimpleNamespace(location=DATASET_DIR)!git clone https://github.com/facebookresearch/sam2.git!wget-O/content/sam2/sam2/configs/train.yaml'https://drive.usercontent.google.com/download?id=11cmbxPPsYqFyWq87tmLgBAQ6OZgEhPG3'%cd./sam2/!pip install-e.[dev]-q!cd./checkpoints&&./download_ckpts.sh🧹 修改数据集文件名
SAM-2.1 训练对文件名格式有要求,这一步把 自定义数据集导出的文件名改成兼容形式。
# 将 自定义数据集文件名改成 SAM 2.1 兼容格式。# 也可以通过调整 sam2/sam2/configs/train.yaml 来避免这一步。importosimportre FOLDER="/content/data/train"forfilenameinos.listdir(FOLDER):# Replace all except last dot with underscorenew_filename=filename.replace(".","_",filename.count(".")-1)ifnotre.search(r"_\d+\.\w+$",new_filename):# Add an int to the end of base namenew_filename=new_filename.replace(".","_1.")os.rename(os.path.join(FOLDER,filename),os.path.join(FOLDER,new_filename))🏋️ 开始训练 SAM-2.1
使用官方训练脚本和配置文件启动单 GPU 训练,并可通过 TensorBoard 查看训练曲线。
!python training/train.py-c'configs/train.yaml'--use-cluster0--num-gpus1%load_ext tensorboard%tensorboard--bind_all--logdir./sam2_logs/🔍 可视化微调模型结果
加载微调 checkpoint 和基础 checkpoint,在验证集图片上生成 mask 并并排对比。
!pip install supervision-qimporttorchfromsam2.build_samimportbuild_sam2fromsam2.automatic_mask_generatorimportSAM2AutomaticMaskGeneratorimportsupervisionassvimportosimportrandomfromPILimportImageimportnumpyasnp# 在整个 notebook 中使用 bfloat16# from Meta notebooktorch.autocast("cuda",dtype=torch.bfloat16).__enter__()iftorch.cuda.get_device_properties(0).major>=8:torch.backends.cuda.matmul.allow_tf32=Truetorch.backends.cudnn.allow_tf32=Truecheckpoint="/content/sam2/sam2_logs/configs/train.yaml/checkpoints/checkpoint.pt"model_cfg="configs/sam2.1/sam2.1_hiera_b+.yaml"sam2=build_sam2(model_cfg,checkpoint,device="cuda")mask_generator=SAM2AutomaticMaskGenerator(sam2)checkpoint_base="/content/sam2/checkpoints/sam2.1_hiera_base_plus.pt"model_cfg_base="configs/sam2.1/sam2.1_hiera_b+.yaml"sam2_base=build_sam2(model_cfg_base,checkpoint_base,device="cuda")mask_generator_base=SAM2AutomaticMaskGenerator(sam2_base)validation_set=os.listdir("/content/data/valid")# choose random with .json extensionimage=random.choice([imgforimginvalidation_setifimg.endswith(".jpg")])image=os.path.join("/content/data/valid",image)opened_image=np.array(Image.open(image).convert("RGB"))result=mask_generator.generate(opened_image)detections=sv.Detections.from_sam(sam_result=result)mask_annotator=sv.MaskAnnotator(color_lookup=sv.ColorLookup.INDEX)annotated_image=opened_image.copy()annotated_image=mask_annotator.annotate(annotated_image,detections=detections)base_annotator=sv.MaskAnnotator(color_lookup=sv.ColorLookup.INDEX)base_result=mask_generator_base.generate(opened_image)base_detections=sv.Detections.from_sam(sam_result=base_result)base_annotated_image=opened_image.copy()base_annotated_image=base_annotator.annotate(base_annotated_image,detections=base_detections)sv.plot_images_grid(images=[annotated_image,base_annotated_image],titles=["Fine-Tuned SAM-2.1","Base SAM-2.1"],grid_size=(1,2))📌 小结
这篇教程完整整理了Fine-Tune SAM-2.1的核心复现流程。实际操作时,建议先确认 GPU、依赖版本、数据集路径和模型权限,再逐段运行 notebook。
- 下载 SAM2 格式数据集
- 克隆并安装 SAM2 官方仓库
- 修正数据集文件名以适配训练配置
- 启动 SAM-2.1 单 GPU 微调
- 对比微调模型和基础模型的自动 mask 结果
后续我会继续按源项目顺序整理 项目教程 中的目标检测、实例分割、OCR、多目标跟踪和视觉大模型教程。
📚 同系列教程汇总
Google Gemini 3.5 Flash 零样本目标检测教程:从提示词到可视化结果
GLM-OCR 文档识别实战教程:从验证码、公式到车牌 OCR
RF-DETR + ByteTrack 多目标跟踪实战教程:从命令行到 Python 视频轨迹可视化
SAM 3 图像分割实战教程:文本、框和点提示的多种分割方式
SAM 3 视频分割实战教程:用文本提示分割并跟踪视频中的目标
SAM-2.1 自定义分割微调实战:从自定义数据集到自动掩码生成-本文