train_segmentation.py 2.0 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354
  1. import paddlex as pdx
  2. from paddlex import transforms as T
  3. # 定义训练和验证时的transforms
  4. # API说明:https://github.com/PaddlePaddle/PaddleX/blob/release/2.0-rc/paddlex/cv/transforms/operators.py
  5. train_transforms = T.Compose([
  6. T.Resize(target_size=512),
  7. T.RandomHorizontalFlip(),
  8. T.Normalize(
  9. mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5]),
  10. ])
  11. eval_transforms = T.Compose([
  12. T.Resize(target_size=512),
  13. T.Normalize(
  14. mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5]),
  15. ])
  16. # 下载和解压指针刻度分割数据集,如果已经预先下载,可注视掉下面两行
  17. meter_seg_dataset = 'https://bj.bcebos.com/paddlex/examples/meter_reader/datasets/meter_seg.tar.gz'
  18. pdx.utils.download_and_decompress(meter_seg_dataset, path='./')
  19. # 定义训练和验证所用的数据集
  20. # API说明:https://github.com/PaddlePaddle/PaddleX/blob/release/2.0-rc/paddlex/cv/datasets/seg_dataset.py#L22
  21. train_dataset = pdx.datasets.SegDataset(
  22. data_dir='meter_seg',
  23. file_list='meter_seg/train.txt',
  24. label_list='meter_seg/labels.txt',
  25. transforms=train_transforms,
  26. shuffle=True)
  27. eval_dataset = pdx.datasets.SegDataset(
  28. data_dir='meter_seg',
  29. file_list='meter_seg/val.txt',
  30. label_list='meter_seg/labels.txt',
  31. transforms=eval_transforms,
  32. shuffle=False)
  33. # 初始化模型,并进行训练
  34. # 可使用VisualDL查看训练指标,参考https://github.com/PaddlePaddle/PaddleX/tree/release/2.0-rc/tutorials/train#visualdl可视化训练指标
  35. num_classes = len(train_dataset.labels)
  36. model = pdx.seg.DeepLabV3P(
  37. num_classes=num_classes, backbone='ResNet50_vd', use_mixed_loss=True)
  38. # API说明:https://github.com/PaddlePaddle/PaddleX/blob/release/2.0-rc/paddlex/cv/models/segmenter.py#L150
  39. # 各参数介绍与调整说明:https://paddlex.readthedocs.io/zh_CN/develop/appendix/parameters.html
  40. model.train(
  41. num_epochs=20,
  42. train_dataset=train_dataset,
  43. train_batch_size=4,
  44. pretrain_weights='IMAGENET',
  45. eval_dataset=eval_dataset,
  46. learning_rate=0.1,
  47. save_dir='output/deeplabv3p_r50vd')