sam2 训练自己的数据集，超详细步骤

yaso_zhang

2519人浏览 · 2025-10-31 11:45:02

yaso_zhang · 2025-10-31 11:45:02 发布

1.拉取sam2的源码（）

git clone https://github.com/facebookresearch/sam2.git

拉取后，该代码需要PyTorch、TorchVisionpython>=3.10和torch>=2.5.1TorchVision 。请按照此处的torchvision>=0.20.1，要安装对应库。

安装后运行

python -c "import torch; print(torch.__version__); print(torch.cuda.is_available()); print(torch.cuda.device_count())"

python -c "import torch, torchvision; print('torch version:', torch.__version__); print('torchvision version:', torchvision.__version__)"

正确会打印版本信息，和是否支持cuda。错误重新安装。conda环境下建议conda直接安装

2按步骤安装依赖：

cd sam2

pip install -e .

3.然后下载训练权重

cd checkpoints && \
./download_ckpts.sh && \
cd ..

4.测试运行sam2

import torch
import numpy as np
from PIL import Image
import matplotlib.pyplot as plt
from sam2.build_sam import build_sam2
from sam2.sam2_image_predictor import SAM2ImagePredictor

checkpoint = "./checkpoints/sam2.1_hiera_large.pt"
model_cfg = "configs/sam2.1/sam2.1_hiera_l.yaml"

# 创建 predictor
predictor = SAM2ImagePredictor(build_sam2(model_cfg, checkpoint))

# 加载图片
image_path = "~/sam2/notebooks/images/cars.jpg" #自己的项目路径
image = np.array(Image.open(image_path).convert("RGB"))
predictor.set_image(image)

# 整图预测，不提供提示点
input_prompts = None

# 执行预测
with torch.inference_mode(), torch.autocast("cuda", dtype=torch.bfloat16):
    masks, scores, logits = predictor.predict(input_prompts)

# 显示第一个 mask
mask = masks[0]  # numpy array
plt.imshow(image)
plt.imshow(mask, alpha=0.5)  # 半透明叠加
plt.show()

5.上面的操作后，下面才是重点。第一下载项目自己的训练集

https://github.com/henghuiding/MOSE-api?tab=readme-ov-file#downloadhttps://github.com/henghuiding/MOSE-api?tab=readme-ov-file#download下载一个小版本的，我下载的是valid.tar.gz。然后解压到文件夹中。

我直接放 ~/sam2trainDir/valid/valid.解压后有JPEGImages和Annotations两个文件夹。里面是原始图片和掩码图片

6.复制或者直接使用项目中的已有配置文件~/sam2/sam2/configs/sam2.1_training/sam2.1_hiera_b+_MOSE_finetune.yaml

这个文件是写好的视频分帧数据训练。我们只要修改其中的配置参数就可以运行。

7 修改参数

修改路径为你解压的文件。

img_folder: /home/bing/sam2trainDir/valid/valid/JPEGImages
gt_folder: /home/bing/sam2trainDir/valid/valid/Annotations
file_list_txt: /home/bing/sam2trainDir/validtest.txt

其中validtest.txt写入你要训练的视频名：比如JPEGImages/002845.可以一次写多个，按列写好就行。然后记得修改scratch:
resolution: 640 #这个是分辨率，太高gpu吃不消
train_batch_size: 1
num_train_workers: 10
num_frames: 1 #这个是每个视频读取的帧数，主要文件中你处理了几帧数据写几，valid只有一帧，记得改成1
max_num_objects: 3 #每帧图像中最多训练或预测的对象数量。
base_lr: 5.0e-6
vision_lr: 3.0e-06
phases_per_epoch: 1 #每个 epoch 的阶段数。
num_epochs: 40 #总训练轮数（epoch）。

详细如下

resolution: 640

输入图像的分辨率（短边或长边？通常是短边 resize 至 640 像素）。

影响训练时图像缩放、内存占用以及模型对小目标的感知能力。

较小分辨率：训练快，占用内存少，但细节损失；

较大分辨率：保留更多细节，但显存占用增加。

train_batch_size: 1

每次训练迭代送入模型的样本数。

你这里是 1，说明 GPU 显存可能有限或者你使用单帧训练。

batch size 太小会导致梯度不稳定，但适合显存受限的情况。

num_train_workers: 10

数据加载器（DataLoader）的工作线程数。

用于并行读取数据，提高训练数据加载速度。

太多线程可能导致 CPU 占用过高或内存竞争。

num_frames: 1

每个样本包含的帧数。

如果是视频分割任务，num_frames 可以大于 1 表示一次输入多帧进行训练。

你这里是 1，说明每次训练只用单帧图像。

max_num_objects: 3

每帧图像中最多训练或预测的对象数量。

超过这个数量的对象可能被忽略或者限制在 3 个实例。

对小目标场景或资源受限训练有用。

base_lr: 5.0e-6

主学习率（base learning rate），控制整个模型参数更新的步长。

很小的学习率说明训练很谨慎，通常用于从头训练或者微调大型模型。

vision_lr: 3.0e-06

图像编码器（Vision Encoder）的学习率。

SAM2 中视觉编码器可能需要比解码器更小的学习率，以防破坏预训练权重。

当训练从零开始（scratch）时，可以设置与 base_lr 相近；微调时通常更小。

phases_per_epoch: 1

每个 epoch 的阶段数。

SAM2 训练可能分成多个阶段（phase），例如 warm-up、fine-tune 等。

这里设为 1，表示每个 epoch 只有一个训练阶段。

num_epochs: 40

总训练轮数（epoch）。

每个 epoch 遍历一次完整训练集。

40 个 epoch 对于小数据集来说可能足够，但大数据集通常需要更多轮。

然后运行训练语句在sam2文件夹下，注意是第一层sam2文件夹，这个文件夹下还有一个sam2文件夹

python -m training.train -c sam2.1_my_hiera_b+_MOSE_finetune.yaml --use-cluster 0 --num-gpus 2

这样就开始训练了。

如果是自己的数据集，记得下载labelme ，进行数据标注。得到的json文件使用python,调用的也是labelme的接口转成同名的图片放在Annotations中做掩码图片。

#!/usr/bin/env python3
import os
import json
import glob
import numpy as np
from PIL import Image
import labelme

# 数据路径
input_dir = "/mnt/hgfs/fileShare/sam2Data/result"  # JSON 文件所在目录
output_dir = os.path.join(input_dir, "masks")      # 保存黑白遮罩的目录
os.makedirs(output_dir, exist_ok=True)

# 遍历所有 JSON 文件
json_files = glob.glob(os.path.join(input_dir, "*.json"))

for json_file in json_files:
    with open(json_file, "r") as f:
        data = json.load(f)
    
    # 原图大小
    img_height = data["imageHeight"]
    img_width = data["imageWidth"]
    img_shape = (img_height, img_width)
    
    # 构造 label_name_to_value，所有标注统一为 1
    label_name_to_value = {shape["label"]: 1 for shape in data["shapes"]}
    
    # 生成 mask
    mask, _ = labelme.utils.shapes_to_label(img_shape, data["shapes"], label_name_to_value)
    
    # 转成黑白 mask
    mask = (mask > 0).astype(np.uint8) * 255
    
    # 保存
    mask_img = Image.fromarray(mask)
    base_name = os.path.splitext(os.path.basename(json_file))[0]
    mask_path = os.path.join(output_dir, f"{base_name}.png")
    mask_img.save(mask_path)
    
    print(f"生成黑白遮罩: {mask_path}")

print("全部完成！")

保证文件格式是这样的

~/sam2trainDir/newDateSet/
├── JPEGImages/
│ ├── 112233/
│ │ ├── 00000.png
│ │ ├── 00001.png
│ │ └── ...
└── Annotations/
├── 112233/
│ ├── 00000.png
│ ├── 00001.png
│ └── ...

这样训练接口调用不报错。

训练结束后会生成新的权重