当前位置：首页 > news >正文

【SAM2代码解析】training部分-1总体概述

news 来源：原创 2025/4/23 9:09:52

总览

1.1 文件总览

training folder保存了训练SAM2的相关代码，该代码允许使用者们用他们自己的数据集（图像、视频或两者一起）去微调SAM2
文件结构如下：

dataset文件夹：保存了包含图像和视频数据集以及数据加载器类及其转换
mode文件夹：包含用于训练、微调的主要模型类（SAM2Train）。SAM2Train继承自SAM2Base模型，并提供用于启用SAM2训练或微调的函数。它还接受所有用于模拟用户提示的训练时参数。
utils：此文件夹包含训练工具，例如日志记录器和分布式训练工具
scripts：此文件夹包含用于提取SA-V数据集帧以用于训练的脚本
loss_fns.py：此文件夹包含训练的主要损失类
optimizer.py：此文件夹包含支持任意调度器的所有优化器工具
trainer.py：此文件包含Trainer类，接收所有的Hydra可配置模块（模型、优化器、数据集等），并实现主要的训练、评估循环。
train.py：此脚本用于启动训练作业，它支持单节点和多节点作业。可运行python training/train.y -h的方式查看有关使用方法。

1.2 训练/微调步骤

以MOSE数据集为例

1、运行pip install -e ".[dev]"安装训练所需的包
2、在configs/sam2.1_training/sam2.1_hiera_b+MOSE_finetune.yaml中设置MOSE数据集的路径

dataset:
#数据集路径
img_folder:null #MOSE JPEGImages 文件夹路径
gt_folder:null #MOSE Annotations 文件径
file_list txt:null # 可选路径，包含用于洲练的视频子集的文件列表

3、使用GPU在MOSE上微调基础模型

python training/train.py -c configs/sam2.1_training/sam2.1_hiera_b+MOSE_finetune.yaml --use-cluster 0

4、可以使用存储在实验日志目录下的tensorboard/文件夹中的TensorBoard日志来监控训练损失。我们还为评估目的提供了一个样本验证拆分。
5、训练完成后，可以使用实验日志目录中checkpoints/文件夹中保存的新检查点
6、在图像和视频上进行训练：代码支持在图像和视频上进行训练，我们提供了用于加载SA-1B作为示例图像数据集、SA-V作为示例视频数据集以及任何DAVIS风格的视频数据集的类。注意在SA-V上进行训练，必须先使用提供的提取脚本将所有视频提取为JPEG帧。

data:
train:
target: training.dataset.sam2_datasets.TorchTrainMixedDataset
phases_per_epoch: ${phases_per_epoch} # 将一个epoch划分成更小的阶段
batch_sizes:
- ${bs1} # 数据集1的批量大小
- ${bs2} # 数据集2的批量大小
datasets:
# SA1B 作为图像数据集的示例
- target: training.dataset.vos_dataset.VOSDataset
training: true
video_dataset:
target: training.dataset.vos_raw_dataset.SA1BRawDataset
img_folder: ${path_to_img_folder}
gt_folder: ${path_to_gt_folder}
file_list_txt: ${path_to_train_filelist} # 可选
sampler:
target: training.dataset.vos_sampler.RandomUniformSampler
num_frames: 1
max_num_objects: ${max_num_objects_per_image}
transforms: ${image_transforms}
# SA-V 作为视频数据集
- target: training.dataset.vos_dataset.VOSDataset
training: true
video_dataset:
target: training.dataset.vos_raw_dataset.JSONRawDataset
img_folder: ${path_to_img_folder}
gt_folder: ${path_to_gt_folder}
file_list_txt: ${path_to_train_filelist} # Optional
ann_every: 4
sampler:
target: training.dataset.vos_sampler.RandomUniformSampler
num_frames: 8 # Number of frames per video
max_num_objects: ${max_num_objects_per_video}
reverse_time_prob: ${reverse_time_prob} # probability to reverse video
transforms: ${video_transforms}
shuffle: True
num_workers: ${num_train_workers}
pin_memory: True
drop_last: True
collate_fn:
target: training.utils.data_utils.collate_fn
partial: true
dict_key: all