faster_rcnn_sar_ship.py 1.8 KB

12345678910111213141516171819202122232425262728293031323334353637383940414243444546474849505152535455565758596061626364
  1. import os
  2. import paddlers as pdrs
  3. from paddlers import transforms as T
  4. # download dataset
  5. data_dir = 'sar_ship_1'
  6. if not os.path.exists(data_dir):
  7. dataset_url = 'https://paddleseg.bj.bcebos.com/dataset/sar_ship_1.tar.gz'
  8. pdrs.utils.download_and_decompress(dataset_url, path='./')
  9. # define transforms
  10. train_transforms = T.Compose([
  11. T.RandomDistort(),
  12. T.RandomExpand(im_padding_value=[123.675, 116.28, 103.53]),
  13. T.RandomCrop(),
  14. T.RandomHorizontalFlip(),
  15. T.BatchRandomResize(
  16. target_sizes=[320, 352, 384, 416, 448, 480, 512, 544, 576, 608],
  17. interp='RANDOM'),
  18. T.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
  19. ])
  20. eval_transforms = T.Compose([
  21. T.Resize(
  22. target_size=608, interp='CUBIC'), T.Normalize(
  23. mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
  24. ])
  25. # define dataset
  26. train_file_list = os.path.join(data_dir, 'train.txt')
  27. val_file_list = os.path.join(data_dir, 'valid.txt')
  28. label_file_list = os.path.join(data_dir, 'labels.txt')
  29. train_dataset = pdrs.datasets.VOCDetection(
  30. data_dir=data_dir,
  31. file_list=train_file_list,
  32. label_list=label_file_list,
  33. transforms=train_transforms,
  34. shuffle=True)
  35. eval_dataset = pdrs.datasets.VOCDetection(
  36. data_dir=data_dir,
  37. file_list=train_file_list,
  38. label_list=label_file_list,
  39. transforms=eval_transforms,
  40. shuffle=False)
  41. # define models
  42. num_classes = len(train_dataset.labels)
  43. model = pdrs.tasks.det.FasterRCNN(num_classes=num_classes)
  44. # train
  45. model.train(
  46. num_epochs=60,
  47. train_dataset=train_dataset,
  48. train_batch_size=2,
  49. eval_dataset=eval_dataset,
  50. pretrain_weights='COCO',
  51. learning_rate=0.005 / 12,
  52. warmup_steps=10,
  53. warmup_start_lr=0.0,
  54. save_interval_epochs=5,
  55. lr_decay_epochs=[20, 40],
  56. save_dir='output/faster_rcnn_sar_ship',
  57. use_vdl=True)