train_cd.py 3.4 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108
  1. #!/usr/bin/env bash
  2. import os.path as osp
  3. import paddle
  4. import paddlers as pdrs
  5. from paddlers import transforms as T
  6. from custom_model import CustomModel
  7. from custom_trainer import make_trainer_and_build
  8. # 数据集路径
  9. DATA_DIR = 'data/levircd/'
  10. # 保存实验结果的路径
  11. EXP_DIR = 'exp/levircd/custom_model/'
  12. # 定义训练和验证时使用的数据变换(数据增强、预处理等)
  13. # 使用Compose组合多种变换方式。Compose中包含的变换将按顺序串行执行
  14. # API说明:https://github.com/PaddlePaddle/PaddleRS/blob/develop/docs/apis/data.md
  15. train_transforms = T.Compose([
  16. # 随机翻转和旋转
  17. T.RandomFlipOrRotate(
  18. # 以0.35的概率执行随机翻转,0.35的概率执行随机旋转
  19. probs=[0.35, 0.35],
  20. # 以0.5的概率执行随机水平翻转,0.5的概率执行随机垂直翻转
  21. probsf=[0.5, 0.5, 0, 0, 0],
  22. # 分别以0.33、0.34和0.33的概率执行90°、180°和270°旋转
  23. probsr=[0.33, 0.34, 0.33]),
  24. # 将数据归一化到[-1,1]
  25. T.Normalize(
  26. mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5])
  27. ])
  28. eval_transforms = T.Compose([
  29. # 验证阶段与训练阶段的数据归一化方式必须相同
  30. T.Normalize(
  31. mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5])
  32. ])
  33. # 分别构建训练、验证和测试所用的数据集
  34. train_dataset = pdrs.datasets.CDDataset(
  35. data_dir=DATA_DIR,
  36. file_list=osp.join(DATA_DIR, 'train.txt'),
  37. label_list=None,
  38. transforms=train_transforms,
  39. num_workers=0,
  40. shuffle=True,
  41. with_seg_labels=False,
  42. binarize_labels=True)
  43. val_dataset = pdrs.datasets.CDDataset(
  44. data_dir=DATA_DIR,
  45. file_list=osp.join(DATA_DIR, 'val.txt'),
  46. label_list=None,
  47. transforms=eval_transforms,
  48. num_workers=0,
  49. shuffle=False,
  50. with_seg_labels=False,
  51. binarize_labels=True)
  52. test_dataset = pdrs.datasets.CDDataset(
  53. data_dir=DATA_DIR,
  54. file_list=osp.join(DATA_DIR, 'test.txt'),
  55. label_list=None,
  56. # 与验证阶段使用相同的数据变换算子
  57. transforms=eval_transforms,
  58. num_workers=0,
  59. shuffle=False,
  60. with_seg_labels=False,
  61. binarize_labels=True)
  62. # 构建自定义模型CustomModel并为其自动生成训练器
  63. # make_trainer_and_build()的首个参数为模型类型,剩余参数为模型构造所需参数
  64. model = make_trainer_and_build(CustomModel, in_channels=3)
  65. # 构建学习率调度器
  66. # 使用定步长学习率衰减策略
  67. lr_scheduler = paddle.optimizer.lr.StepDecay(
  68. learning_rate=0.002, step_size=35000, gamma=0.2)
  69. # 构建优化器
  70. optimizer = paddle.optimizer.Adam(
  71. parameters=model.net.parameters(), learning_rate=lr_scheduler)
  72. # 执行模型训练
  73. model.train(
  74. num_epochs=50,
  75. train_dataset=train_dataset,
  76. train_batch_size=8,
  77. eval_dataset=val_dataset,
  78. optimizer=optimizer,
  79. # 每多少个epoch验证并保存一次模型
  80. save_interval_epochs=5,
  81. # 每多少次迭代记录一次日志
  82. log_interval_steps=50,
  83. save_dir=EXP_DIR,
  84. # 是否使用early stopping策略,当精度不再改善时提前终止训练
  85. early_stop=False,
  86. # 是否启用VisualDL日志功能
  87. use_vdl=True,
  88. # 指定从某个检查点继续训练
  89. resume_checkpoint=None)
  90. # 加载验证集上效果最好的模型
  91. model = pdrs.tasks.load_model(osp.join(EXP_DIR, 'best_model'))
  92. # 在测试集上计算精度指标
  93. res = model.evaluate(test_dataset)
  94. print(res)