ppyolo.py 2.5 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869
  1. # 环境变量配置,用于控制是否使用GPU
  2. # 说明文档:https://paddlex.readthedocs.io/zh_CN/develop/appendix/parameters.html#gpu
  3. import os
  4. os.environ['CUDA_VISIBLE_DEVICES'] = '0'
  5. from paddlex.det import transforms
  6. import paddlex as pdx
  7. # 下载和解压昆虫检测数据集
  8. insect_dataset = 'https://bj.bcebos.com/paddlex/datasets/insect_det.tar.gz'
  9. pdx.utils.download_and_decompress(insect_dataset, path='./')
  10. # 定义训练和验证时的transforms
  11. # API说明 https://paddlex.readthedocs.io/zh_CN/develop/apis/transforms/det_transforms.html
  12. train_transforms = transforms.Compose([
  13. transforms.MixupImage(mixup_epoch=250), transforms.RandomDistort(),
  14. transforms.RandomExpand(), transforms.RandomCrop(), transforms.Resize(
  15. target_size=608, interp='RANDOM'), transforms.RandomHorizontalFlip(),
  16. transforms.Normalize()
  17. ])
  18. eval_transforms = transforms.Compose([
  19. transforms.Resize(
  20. target_size=608, interp='CUBIC'), transforms.Normalize()
  21. ])
  22. # 定义训练和验证所用的数据集
  23. # API说明:https://paddlex.readthedocs.io/zh_CN/develop/apis/datasets.html#paddlex-datasets-vocdetection
  24. train_dataset = pdx.datasets.VOCDetection(
  25. data_dir='insect_det',
  26. file_list='insect_det/train_list.txt',
  27. label_list='insect_det/labels.txt',
  28. transforms=train_transforms,
  29. shuffle=True)
  30. eval_dataset = pdx.datasets.VOCDetection(
  31. data_dir='insect_det',
  32. file_list='insect_det/val_list.txt',
  33. label_list='insect_det/labels.txt',
  34. transforms=eval_transforms)
  35. # 初始化模型,并进行训练
  36. # 可使用VisualDL查看训练指标,参考https://paddlex.readthedocs.io/zh_CN/develop/train/visualdl.html
  37. num_classes = len(train_dataset.labels)
  38. # API说明: https://paddlex.readthedocs.io/zh_CN/develop/apis/models/detection.html#paddlex-det-yolov3
  39. model = pdx.det.YOLOv3(
  40. num_classes=num_classes,
  41. backbone='ResNet50_vd',
  42. with_dcn_v2=True,
  43. use_coord_conv=True,
  44. use_iou_aware=True,
  45. use_spp=True,
  46. use_drop_block=True,
  47. scale_x_y=1.05,
  48. use_iou_loss=True,
  49. use_matrix_nms=True)
  50. # API说明: https://paddlex.readthedocs.io/zh_CN/develop/apis/models/detection.html#train
  51. # 各参数介绍与调整说明:https://paddlex.readthedocs.io/zh_CN/develop/appendix/parameters.html
  52. model.train(
  53. num_epochs=270,
  54. train_dataset=train_dataset,
  55. train_batch_size=8,
  56. eval_dataset=eval_dataset,
  57. learning_rate=0.000125,
  58. lr_decay_epochs=[210, 240],
  59. use_ema=True,
  60. save_dir='output/ppyolo',
  61. use_vdl=True)