farseg_test.py 2.0 KB

12345678910111213141516171819202122232425262728293031323334353637383940414243444546474849505152535455565758
  1. import os
  2. os.environ['CUDA_VISIBLE_DEVICES'] = '0'
  3. import paddlers as pdrs
  4. from paddlers import transforms as T
  5. # 下载和解压视盘分割数据集
  6. optic_dataset = 'https://bj.bcebos.com/paddlex/datasets/optic_disc_seg.tar.gz'
  7. pdrs.utils.download_and_decompress(optic_dataset, path='./')
  8. # 定义训练和验证时的transforms
  9. # API说明:https://github.com/PaddlePaddle/paddlers/blob/develop/docs/apis/transforms/transforms.md
  10. train_transforms = T.Compose([
  11. T.Resize(target_size=512),
  12. T.RandomHorizontalFlip(),
  13. T.Normalize(
  14. mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5]),
  15. ])
  16. eval_transforms = T.Compose([
  17. T.Resize(target_size=512),
  18. T.Normalize(
  19. mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5]),
  20. ])
  21. # 定义训练和验证所用的数据集
  22. # API说明:https://github.com/PaddlePaddle/paddlers/blob/develop/docs/apis/datasets.md
  23. train_dataset = pdrs.datasets.SegDataset(
  24. data_dir='optic_disc_seg',
  25. file_list='optic_disc_seg/train_list.txt',
  26. label_list='optic_disc_seg/labels.txt',
  27. transforms=train_transforms,
  28. num_workers=0,
  29. shuffle=True)
  30. eval_dataset = pdrs.datasets.SegDataset(
  31. data_dir='optic_disc_seg',
  32. file_list='optic_disc_seg/val_list.txt',
  33. label_list='optic_disc_seg/labels.txt',
  34. transforms=eval_transforms,
  35. num_workers=0,
  36. shuffle=False)
  37. # 初始化模型,并进行训练
  38. # 可使用VisualDL查看训练指标,参考https://github.com/PaddlePaddle/paddlers/blob/develop/docs/visualdl.md
  39. num_classes = len(train_dataset.labels)
  40. model = pdrs.tasks.FarSeg(num_classes=num_classes)
  41. # API说明:https://github.com/PaddlePaddle/paddlers/blob/develop/docs/apis/models/semantic_segmentation.md
  42. # 各参数介绍与调整说明:https://github.com/PaddlePaddle/paddlers/blob/develop/docs/parameters.md
  43. model.train(
  44. num_epochs=10,
  45. train_dataset=train_dataset,
  46. train_batch_size=4,
  47. eval_dataset=eval_dataset,
  48. learning_rate=0.01,
  49. pretrain_weights=None,
  50. save_dir='output/farseg')