ch_PP-OCRv4_rec_distill.yml

Global:
  device: gpu
  epoch_num: 200
  log_smooth_window: 20
  print_batch_step: 10
  output_dir: ./output/rec/rec_ppocr_v4_distill
  eval_epoch_step: [0, 1]
  cal_metric_during_train: true
  pretrained_model:
  checkpoints:
  use_tensorboard: false
  infer_mode: false
  infer_img: doc/imgs_words/ch/word_1.jpg
  character_dict_path: &character_dict_path torchocr/utils/ppocr_keys_v1.txt
  max_text_length: &max_text_length 25
  use_space_char: &use_space_char true

Export:
  export_dir:
  export_shape: [ 1, 3, 48, 320 ]
  dynamic_axes: [ 0, 2, 3 ]

Optimizer:
  name: Adam
  lr: 0.001
  weight_decay: 3.0e-05

LRScheduler:
  name: CosineAnnealingLR
  warmup_epoch: 5

Architecture:
  model_type: rec
  name: DistillationModel
  algorithm: Distillation
  Models:
    Teacher:
      pretrained:
      freeze_params: true
      return_all_feats: true
      model_type: rec
      algorithm: SVTR
      Transform: null
      Backbone:
        name: PPHGNet_small
      Head:
        name: MultiHead
        head_list:
          - CTCHead:
              Neck:
                name: svtr
                dims: 120
                depth: 2
                hidden_dims: 120
                kernel_size: [1, 3]
                use_guide: True
              Head:
                fc_decay: 0.00001
          - NRTRHead:
              nrtr_dim: 384
              max_text_length: *max_text_length
    Student:
      pretrained:
      freeze_params: false
      return_all_feats: true
      model_type: rec
      algorithm: SVTR
      Transform: null
      Backbone:
        name: PPLCNetV3
        scale: 0.95
      Head:
        name: MultiHead
        head_list:
          - CTCHead:
              Neck:
                name: svtr
                dims: 120
                depth: 2
                hidden_dims: 120
                kernel_size: [1, 3]
                use_guide: True
          - NRTRHead:
              nrtr_dim: 384
              max_text_length: *max_text_length
Loss:
  name: CombinedLoss
  loss_config_list:
  - DistillationDKDLoss:
      weight: 0.1
      model_name_pairs:
      - - Student
        - Teacher
      key: head_out
      multi_head: true
      alpha: 1.0
      beta: 2.0
      dis_head: nrtr
      name: dkd
  - DistillationCTCLoss:
      weight: 1.0
      model_name_list:
      - Student
      key: head_out
      multi_head: true
  - DistillCTCLogits:
      weight: 1.0
      reduction: mean
      model_name_pairs:
      - - Student
        - Teacher
      key: head_out
  - DistillationNRTRLoss:
      weight: 1.0
      smoothing: false
      model_name_list:
      - Student
      key: head_out
      multi_head: true
  
PostProcess:
  name: DistillationCTCLabelDecode
  model_name:
  - Student
  key: head_out
  multi_head: true
  character_dict_path: *character_dict_path
  use_space_char: *use_space_char
  
Metric:
  name: DistillationMetric
  base_metric_name: RecMetric
  main_indicator: acc
  key: Student
  ignore_space: false
Train:
  dataset:
    name: MultiScaleDataSet
    ds_width: false
    data_dir: ./train_data/
    ext_op_transform_idx: 1
    label_file_list:
    - ./train_data/train_list.txt
    transforms:
    - DecodeImage:
        img_mode: BGR
        channel_first: false
    - RecConAug:
        prob: 0.5
        ext_data_num: 2
        image_shape: [48, 320, 3]
        max_text_length: *max_text_length
    - RecAug:
    - MultiLabelEncode:
        gtc_encode: NRTRLabelEncode
    - KeepKeys:
        keep_keys:
        - image
        - label_ctc
        - label_gtc
        - length
        - valid_ratio
  sampler:
    name: MultiScaleSampler
    scales: [[320, 32], [320, 48], [320, 64]]
    first_bs: &bs 192
    fix_bs: false
    divided_factor: [8, 16] # w, h
    is_training: True
  loader:
    shuffle: true
    batch_size_per_card: *bs
    drop_last: true
    num_workers: 8
Eval:
  dataset:
    name: SimpleDataSet
    data_dir: ./train_data
    label_file_list:
    - ./train_data/val_list.txt
    transforms:
    - DecodeImage:
        img_mode: BGR
        channel_first: false
    - MultiLabelEncode:
        gtc_encode: NRTRLabelEncode
    - RecResizeImg:
        image_shape: [3, 48, 320]
    - KeepKeys:
        keep_keys:
        - image
        - label_ctc
        - label_gtc
        - length
        - valid_ratio
  loader:
    shuffle: false
    drop_last: false
    batch_size_per_card: 128
    num_workers: 4