ForensicHub 文档ForensicHub 文档
  • 基本信息

    • 简介
    • 框架设计
  • 快速上手

    • 安装
    • 组件注册
    • Yaml配置
  • 模型

    • 模型汇总
  • 指标

    • 评估指标
  • English
  • 简体中文
GitHub
  • 基本信息

    • 简介
    • 框架设计
  • 快速上手

    • 安装
    • 组件注册
    • Yaml配置
  • 模型

    • 模型汇总
  • 指标

    • 评估指标
  • English
  • 简体中文
GitHub
  • 指南

    • 简介
    • 框架设计
  • 快速上手

    • 安装
    • 组件注册
    • Yaml配置

Yaml配置

ForensicHub中,用户直接通过Yaml文件配置训练或测试的 pipeline,下面将解释Yaml中具体用到的参数,最后会给出完整的Yaml文件。以下示例的Yaml为仓库首页Readme中的Quick Start中的示例,希望你已经阅读过仓库首页的Readme:

# DDP
gpus: "4,5"
flag: train

gpus定义了多卡环境下使用的显卡序号,示例中使用的是第4,5两张卡。flag标识了训练(train)和测试(test)阶段。

# Log
log_dir: "./log/aigc_resnet_df_train"

log_dir定义了存放log和checkpoint的位置。

# Model
model:
  name: Resnet50
  # Model specific setting
  init_config:
    pretrained: true
    num_classes: 1

name必须是注册时使用的名称,否则会找不到注册类。init_config中可以加入模型初始化所使用的参数,具体参数名称要与模型定义的初始化参数的名称相匹配。

# Train dataset
train_dataset:
  name: AIGCLabelDataset
  dataset_name: DiffusionForensics_train
  init_config:
    image_size: 224
    path: /mnt/data1/public_datasets/AIGC/DiffusionForensics/images/train.json
#  Test dataset (one or many)
test_dataset:
  - name: AIGCLabelDataset
    dataset_name: DiffusionForensics_val
    init_config:
      image_size: 224
      path: /mnt/data1/public_datasets/AIGC/DiffusionForensics/images/val.json

train_dataset和test_dataset定义了训练中使用的训练集和测试集,其中test_dataset可以使用一个或多个测试集,使用多个时为列表形式。同样的,name必须是注册时使用的名称,init_config中可以加入初始化所使用的参数。

# Transform
transform:
  name: AIGCTransform

# Evaluators
evaluator:
  - name: ImageF1
    init_config:
      threshold: 0.5

transform和evaluator参数和前面的模型和数据集参数格式相同。

# Training related
batch_size: 768
test_batch_size: 128
epochs: 20
accum_iter: 1
record_epoch: 0  # Save the best only after record epoch.

定义训练使用的参数,其中record_epoch定义了什么时候开始保存性能最好的checkpoint。

# Test related
no_model_eval: false
test_period: 1

no_model_eval控制在测试时是否使用torch中的model_eval,test_period控制多少个epoch后进行一次测试。

# Logging & TensorBoard
log_per_epoch_count: 20

# DDP & AMP settings
find_unused_parameters: false
use_amp: true

log_per_epoch_count控制一个epoch中记录多少次到tensorboard中。

# Optimizer parameters
weight_decay: 0.05
lr: 1e-4
blr: 0.001
min_lr: 1e-5
warmup_epochs: 1
# Device and training control
device: "cuda"
seed: 42
resume: ""
start_epoch: 0
num_workers: 8
pin_mem: true

# Distributed training parameters
world_size: 1
local_rank: -1
dist_on_itp: false
dist_url: "env://"

以下为完整的Yaml文件:

# DDP
gpus: "4,5"
flag: train

# Log
log_dir: "./log/aigc_resnet_df_train"

# Task
if_predict_label: true
if_predict_mask: false

# Model
model:
  name: Resnet50
  # Model specific setting
  init_config:
    pretrained: true
    num_classes: 1

# Train dataset
train_dataset:
  name: AIGCLabelDataset
  dataset_name: DiffusionForensics_train
  init_config:
    image_size: 224
    path: /mnt/data1/public_datasets/AIGC/DiffusionForensics/images/train.json
#  Test dataset (one or many)
test_dataset:
  - name: AIGCLabelDataset
    dataset_name: DiffusionForensics_val
    init_config:
      image_size: 224
      path: /mnt/data1/public_datasets/AIGC/DiffusionForensics/images/val.json

# Transform
transform:
  name: AIGCTransform

# Evaluators
evaluator:
  - name: ImageF1
    init_config:
      threshold: 0.5

# Training related
batch_size: 768
test_batch_size: 128
epochs: 20
accum_iter: 1
record_epoch: 0  # Save the best only after record epoch.

# Test related
no_model_eval: false
test_period: 1

# Logging & TensorBoard
log_per_epoch_count: 20

# DDP & AMP settings
find_unused_parameters: false
use_amp: true

# Optimizer parameters
weight_decay: 0.05
lr: 1e-4
blr: 0.001
min_lr: 1e-5
warmup_epochs: 1

# Device and training control
device: "cuda"
seed: 42
resume: ""
start_epoch: 0
num_workers: 8
pin_mem: true

# Distributed training parameters
world_size: 1
local_rank: -1
dist_on_itp: false
dist_url: "env://"
在 GitHub 上编辑此页
上次更新: 2025/6/16 09:08
贡献者: Bo Du
Prev
组件注册