|
@@ -0,0 +1,109 @@
|
|
|
+#!/usr/bin/env python
|
|
|
+
|
|
|
+# 图像复原模型NAFNet训练示例脚本
|
|
|
+# 执行此脚本前,请确认已正确安装PaddleRS库
|
|
|
+
|
|
|
+import paddle
|
|
|
+
|
|
|
+import paddlers as pdrs
|
|
|
+from paddlers import transforms as T
|
|
|
+
|
|
|
+# 数据集存放目录
|
|
|
+DATA_DIR = './data/RICE1'
|
|
|
+# 训练集`file_list`文件路径
|
|
|
+TRAIN_FILE_LIST_PATH = './data/RICE1/train.txt'
|
|
|
+# 验证集`file_list`文件路径
|
|
|
+EVAL_FILE_LIST_PATH = './data/RICE1/val.txt'
|
|
|
+# 实验目录,保存输出的模型权重和结果
|
|
|
+EXP_DIR = './output/nafnet/'
|
|
|
+
|
|
|
+# 下载和解压遥感影像去云数据集
|
|
|
+pdrs.utils.download_and_decompress(
|
|
|
+ 'https://paddlers.bj.bcebos.com/datasets/RICE1.zip', path='./data/')
|
|
|
+
|
|
|
+# 定义训练和验证时使用的数据变换(数据增强、预处理等)
|
|
|
+# 使用Compose组合多种变换方式。Compose中包含的变换将按顺序串行执行
|
|
|
+# API说明:https://github.com/PaddlePaddle/PaddleRS/blob/develop/docs/apis/data.md
|
|
|
+train_transforms = [
|
|
|
+ # 从输入影像中裁剪256×256大小的影像块
|
|
|
+ T.RandomCrop(crop_size=256),
|
|
|
+ # 以50%的概率实施随机水平翻转
|
|
|
+ T.RandomHorizontalFlip(prob=0.5),
|
|
|
+ # 以50%的概率实施随机垂直翻转
|
|
|
+ T.RandomVerticalFlip(prob=0.5),
|
|
|
+ # 以默认设置实施随机的翻转或旋转
|
|
|
+ T.RandomFlipOrRotate(),
|
|
|
+ # 将数据归一化到[0,1]
|
|
|
+ T.Normalize(
|
|
|
+ mean=[0.0, 0.0, 0.0], std=[1.0, 1.0, 1.0])
|
|
|
+]
|
|
|
+
|
|
|
+eval_transforms = [
|
|
|
+ # 验证阶段与训练阶段的数据归一化方式必须相同
|
|
|
+ T.Normalize(
|
|
|
+ mean=[0.0, 0.0, 0.0], std=[1.0, 1.0, 1.0])
|
|
|
+]
|
|
|
+
|
|
|
+# 分别构建训练和验证所用的数据集
|
|
|
+train_dataset = pdrs.datasets.ResDataset(
|
|
|
+ data_dir=DATA_DIR,
|
|
|
+ file_list=TRAIN_FILE_LIST_PATH,
|
|
|
+ transforms=train_transforms,
|
|
|
+ num_workers=0,
|
|
|
+ shuffle=True,
|
|
|
+ sr_factor=None)
|
|
|
+
|
|
|
+eval_dataset = pdrs.datasets.ResDataset(
|
|
|
+ data_dir=DATA_DIR,
|
|
|
+ file_list=EVAL_FILE_LIST_PATH,
|
|
|
+ transforms=eval_transforms,
|
|
|
+ num_workers=0,
|
|
|
+ shuffle=False,
|
|
|
+ sr_factor=None)
|
|
|
+
|
|
|
+# 使用以下参数构建NAFNet模型
|
|
|
+in_channels = 3
|
|
|
+width = 32
|
|
|
+middle_blk_num = 12
|
|
|
+enc_blk_nums = [2, 2, 4, 8]
|
|
|
+dec_blk_nums = [2, 2, 2, 2]
|
|
|
+# 目前已支持的模型请参考:https://github.com/PaddlePaddle/PaddleRS/blob/develop/docs/intro/model_zoo.md
|
|
|
+# 模型输入参数请参考:https://github.com/PaddlePaddle/PaddleRS/blob/develop/paddlers/tasks/restorer.py
|
|
|
+
|
|
|
+model = pdrs.tasks.res.NAFNet(
|
|
|
+ in_channels=in_channels,
|
|
|
+ width=width,
|
|
|
+ middle_blk_num=middle_blk_num,
|
|
|
+ enc_blk_nums=enc_blk_nums,
|
|
|
+ dec_blk_nums=dec_blk_nums)
|
|
|
+
|
|
|
+# 制定余弦学习率衰减策略
|
|
|
+lr_scheduler = paddle.optimizer.lr.CosineAnnealingDecay(
|
|
|
+ learning_rate=0.0006, T_max=4000, eta_min=8e-7)
|
|
|
+
|
|
|
+# 构造AdamW优化器
|
|
|
+optimizer = paddle.optimizer.AdamW(
|
|
|
+ learning_rate=lr_scheduler,
|
|
|
+ parameters=model.net.parameters(),
|
|
|
+ weight_decay=0.0,
|
|
|
+ beta1=0.9,
|
|
|
+ beta2=0.9,
|
|
|
+ epsilon=1e-8)
|
|
|
+
|
|
|
+# 执行模型训练
|
|
|
+model.train(
|
|
|
+ num_epochs=200,
|
|
|
+ train_dataset=train_dataset,
|
|
|
+ train_batch_size=20,
|
|
|
+ eval_dataset=eval_dataset,
|
|
|
+ optimizer=optimizer,
|
|
|
+ save_interval_epochs=10,
|
|
|
+ # 每多少次迭代记录一次日志
|
|
|
+ log_interval_steps=10,
|
|
|
+ save_dir=EXP_DIR,
|
|
|
+ # 是否使用early stopping策略,当精度不再改善时提前终止训练
|
|
|
+ early_stop=False,
|
|
|
+ # 是否启用VisualDL日志功能
|
|
|
+ use_vdl=True,
|
|
|
+ # 指定从某个检查点继续训练
|
|
|
+ resume_checkpoint=None)
|