finetune_unet.md 3.3 KB
Newer Older
Z
Zeyu Chen 已提交
1
# U-Net模型训练教程
W
wuzewu 已提交
2

Z
Zeyu Chen 已提交
3
* 本教程旨在介绍如何通过使用PaddleSeg提供的 ***`U-Net`*** 预训练模型在自定义数据集上进行训练
W
wuzewu 已提交
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

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

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

## 一. 准备待训练数据

我们提前准备好了一份数据集,通过以下代码进行下载

```shell
python dataset/download_pet.py
```

## 二. 下载预训练模型

关于PaddleSeg支持的所有预训练模型的列表,我们可以从[模型组合](#模型组合)中查看我们所需模型的名字和配置。

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

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

## 三. 准备配置

W
wuzewu 已提交
29
接着我们需要确定相关配置,从本教程的角度,配置分为三部分:
W
wuzewu 已提交
30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49

* 数据集
  * 训练集主目录
  * 训练集文件列表
  * 测试集文件列表
  * 评估集文件列表
* 预训练模型
  * 预训练模型名称
  * 预训练模型的backbone网络
  * 预训练模型的Normalization类型
  * 预训练模型路径
* 其他
  * 学习率
  * Batch大小
  * ...

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

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

W
wuyefeilin 已提交
50
其他配置则根据数据集和机器环境的情况进行调节,最终我们保存一个如下内容的yaml配置文件,存放路径为**configs/unet_pet.yaml**
W
wuzewu 已提交
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

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


# 预训练模型配置
MODEL:
    MODEL_NAME: "unet"
    DEFAULT_NORM_TYPE: "bn"

# 其他配置
TRAIN_CROP_SIZE: (512, 512)
EVAL_CROP_SIZE: (512, 512)
AUG:
    AUG_METHOD: "unpadding"
    FIX_RESIZE_SIZE: (512, 512)
BATCH_SIZE: 4
TRAIN:
W
wuzewu 已提交
76
    PRETRAINED_MODEL_DIR: "./pretrained_model/unet_bn_coco/"
Z
Zeyu Chen 已提交
77
    MODEL_SAVE_DIR: "./saved_model/unet_pet/"
W
wuzewu 已提交
78 79
    SNAPSHOT_EPOCH: 10
TEST:
Z
Zeyu Chen 已提交
80
    TEST_MODEL: "./saved_model/unet_pet/final"
W
wuzewu 已提交
81
SOLVER:
Z
Zeyu Chen 已提交
82
    NUM_EPOCHS: 100
W
wuzewu 已提交
83 84 85 86 87
    LR: 0.005
    LR_POLICY: "poly"
    OPTIMIZER: "adam"
```

W
wuzewu 已提交
88
## 四. 配置/数据校验
W
wuzewu 已提交
89

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

```shell
W
wuyefeilin 已提交
93
python pdseg/check.py --cfg ./configs/unet_pet.yaml
W
wuzewu 已提交
94 95 96 97 98 99
```


## 五. 开始训练

校验通过后,使用下述命令启动训练
W
wuzewu 已提交
100 101

```shell
W
wuyefeilin 已提交
102
python pdseg/train.py --use_gpu --cfg ./configs/unet_pet.yaml
W
wuzewu 已提交
103 104
```

W
wuzewu 已提交
105
## 六. 进行评估
W
wuzewu 已提交
106

W
wuzewu 已提交
107
模型训练完成,使用下述命令启动评估
W
wuzewu 已提交
108 109

```shell
W
wuyefeilin 已提交
110
python pdseg/eval.py --use_gpu --cfg ./configs/unet_pet.yaml
W
wuzewu 已提交
111 112 113 114 115 116
```

## 模型组合

|预训练模型名称|BackBone|Norm|数据集|配置|
|-|-|-|-|-|
W
wuzewu 已提交
117
|unet_bn_coco|-|bn|COCO|MODEL.MODEL_NAME: unet <br> MODEL.DEFAULT_NORM_TYPE: bn|