finetune_ocrnet.md 4.4 KB
Newer Older
W
wuzewu 已提交
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145
# OCRNet模型使用教程

本教程旨在介绍如何通过使用PaddleSeg提供的 ***`OCRNet`*** 预训练模型在自定义数据集上进行训练、评估和可视化。

* 在阅读本教程前,请确保您已经了解过PaddleSeg的[快速入门](../README.md#快速入门)[基础功能](../README.md#基础功能)等章节,以便对PaddleSeg有一定的了解。

* 本教程的所有命令都基于PaddleSeg主目录进行执行。

* 目前OCRNet只支持HRNet作为backbone

## 一. 准备待训练数据

![](./imgs/optic.png)

我们提前准备好了一份眼底医疗分割数据集,包含267张训练图片、76张验证图片、38张测试图片。通过以下命令进行下载:

```shell
python dataset/download_optic.py
```


## 二. 下载预训练模型

接着下载对应的预训练模型

```shell
python pretrained_model/download_model.py ocrnet_w18_bn_cityscapes
```

关于已有的OCRNet预训练模型的列表,请参见[模型组合](#模型组合)。如果需要使用其他预训练模型,下载该模型并将配置中的BACKBONE、NORM_TYPE等进行替换即可。

## 三. 准备配置

接着我们需要确定相关配置,从本教程的角度,配置分为三部分:

* 数据集
  * 训练集主目录
  * 训练集文件列表
  * 测试集文件列表
  * 评估集文件列表
* 预训练模型
  * 预训练模型名称
  * 预训练模型各阶段通道数设置
  * 预训练模型的Normalization类型
  * 预训练模型路径
* 其他
  * 学习率
  * Batch大小
  * ...

在三者中,预训练模型的配置尤为重要,如果模型配置错误,会导致预训练的参数没有加载,进而影响收敛速度。预训练模型相关的配置如第二步所展示。

数据集的配置和数据路径有关,在本教程中,数据存放在`dataset/optic_disc_seg`

其他配置则根据数据集和机器环境的情况进行调节,最终我们保存一个如下内容的yaml配置文件,存放路径为**configs/ocrnet_optic.yaml**

```yaml
# 数据集配置
DATASET:
    DATA_DIR: "./dataset/optic_disc_seg/"
    NUM_CLASSES: 2
    TEST_FILE_LIST: "./dataset/optic_disc_seg/test_list.txt"
    TRAIN_FILE_LIST: "./dataset/optic_disc_seg/train_list.txt"
    VAL_FILE_LIST: "./dataset/optic_disc_seg/val_list.txt"
    VIS_FILE_LIST: "./dataset/optic_disc_seg/test_list.txt"

# 预训练模型配置
MODEL:
    MODEL_NAME: "ocrnet"
    DEFAULT_NORM_TYPE: "bn"
    # 此处设置backbone的配置
    HRNET:
        STAGE2:
            NUM_CHANNELS: [18, 36]
        STAGE3:
            NUM_CHANNELS: [18, 36, 72]
        STAGE4:
            NUM_CHANNELS: [18, 36, 72, 144]
    # 此处设置OCR HEAD
    

# 其他配置
TRAIN_CROP_SIZE: (512, 512)
EVAL_CROP_SIZE: (512, 512)
AUG:
    AUG_METHOD: "unpadding"
    FIX_RESIZE_SIZE: (512, 512)
BATCH_SIZE: 4
TRAIN:
    PRETRAINED_MODEL_DIR: "./pretrained_model/ocrnet_w18_bn_cityscapes/"
    MODEL_SAVE_DIR: "./saved_model/ocrnet_optic/"
    SNAPSHOT_EPOCH: 5
TEST:
    TEST_MODEL: "./saved_model/ocrnet_optic/final"
SOLVER:
    NUM_EPOCHS: 10
    LR: 0.001
    LR_POLICY: "poly"
    OPTIMIZER: "adam"
```

## 四. 配置/数据校验

在开始训练和评估之前,我们还需要对配置和数据进行一次校验,确保数据和配置是正确的。使用下述命令启动校验流程

```shell
python pdseg/check.py --cfg ./configs/ocrnet_optic.yaml
```


## 五. 开始训练

校验通过后,使用下述命令启动训练

```shell
# 指定GPU卡号(以0号卡为例)
export CUDA_VISIBLE_DEVICES=0
# 训练
python pdseg/train.py --use_gpu --cfg ./configs/ocrnet_optic.yaml
```

## 六. 进行评估

模型训练完成,使用下述命令启动评估

```shell
python pdseg/eval.py --use_gpu --cfg ./configs/ocrnet_optic.yaml
```

## 七. 进行可视化
使用下述命令启动预测和可视化

```shell
python pdseg/vis.py --use_gpu --cfg ./configs/ocrnet_optic.yaml
```

预测结果将保存在visual目录下,以下展示其中1张图片的预测效果:

![](imgs/optic_hrnet.png)

## 模型组合

|预训练模型名称|Backbone|数据集|配置|
|-|-|-|-|
|ocrnet_w18_bn_cityscapes|OCRNet| Cityscapes | MODEL.MODEL_NAME: ocrnet <br> MODEL.HRNET.STAGE2.NUM_CHANNELS: [18, 36] <br> MODEL.HRNET.STAGE3.NUM_CHANNELS: [18, 36, 72] <br> MODEL.HRNET.STAGE4.NUM_CHANNELS: [18, 36, 72, 144] <br> MODEL.DEFAULT_NORM_TYPE: bn <br> MODEL.OCR.OCR_MID_CHANNELS: 512 <br> MODEL.OCR.OCR_KEY_CHANNELS: 256|