pretrain_weights.py 11 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230
  1. import paddlex
  2. import paddlex.utils.logging as logging
  3. import paddlehub as hub
  4. import os
  5. import os.path as osp
  6. image_pretrain = {
  7. 'ResNet18':
  8. 'https://paddle-imagenet-models-name.bj.bcebos.com/ResNet18_pretrained.tar',
  9. 'ResNet34':
  10. 'https://paddle-imagenet-models-name.bj.bcebos.com/ResNet34_pretrained.tar',
  11. 'ResNet50':
  12. 'http://paddle-imagenet-models-name.bj.bcebos.com/ResNet50_pretrained.tar',
  13. 'ResNet101':
  14. 'http://paddle-imagenet-models-name.bj.bcebos.com/ResNet101_pretrained.tar',
  15. 'ResNet50_vd':
  16. 'https://paddle-imagenet-models-name.bj.bcebos.com/ResNet50_vd_pretrained.tar',
  17. 'ResNet101_vd':
  18. 'https://paddle-imagenet-models-name.bj.bcebos.com/ResNet101_vd_pretrained.tar',
  19. 'ResNet50_vd_ssld':
  20. 'https://paddle-imagenet-models-name.bj.bcebos.com/ResNet50_vd_ssld_pretrained.tar',
  21. 'ResNet101_vd_ssld':
  22. 'https://paddle-imagenet-models-name.bj.bcebos.com/ResNet101_vd_ssld_pretrained.tar',
  23. 'MobileNetV1':
  24. 'http://paddle-imagenet-models-name.bj.bcebos.com/MobileNetV1_pretrained.tar',
  25. 'MobileNetV2_x1.0':
  26. 'https://paddle-imagenet-models-name.bj.bcebos.com/MobileNetV2_pretrained.tar',
  27. 'MobileNetV2_x0.5':
  28. 'https://paddle-imagenet-models-name.bj.bcebos.com/MobileNetV2_x0_5_pretrained.tar',
  29. 'MobileNetV2_x2.0':
  30. 'https://paddle-imagenet-models-name.bj.bcebos.com/MobileNetV2_x2_0_pretrained.tar',
  31. 'MobileNetV2_x0.25':
  32. 'https://paddle-imagenet-models-name.bj.bcebos.com/MobileNetV2_x0_25_pretrained.tar',
  33. 'MobileNetV2_x1.5':
  34. 'https://paddle-imagenet-models-name.bj.bcebos.com/MobileNetV2_x1_5_pretrained.tar',
  35. 'MobileNetV3_small':
  36. 'https://paddle-imagenet-models-name.bj.bcebos.com/MobileNetV3_small_x1_0_pretrained.tar',
  37. 'MobileNetV3_large':
  38. 'https://paddle-imagenet-models-name.bj.bcebos.com/MobileNetV3_large_x1_0_pretrained.tar',
  39. 'MobileNetV3_small_x1_0_ssld':
  40. 'https://paddle-imagenet-models-name.bj.bcebos.com/MobileNetV3_small_x1_0_ssld_pretrained.tar',
  41. 'MobileNetV3_large_x1_0_ssld':
  42. 'https://paddle-imagenet-models-name.bj.bcebos.com/MobileNetV3_large_x1_0_ssld_pretrained.tar',
  43. 'DarkNet53':
  44. 'https://paddle-imagenet-models-name.bj.bcebos.com/DarkNet53_ImageNet1k_pretrained.tar',
  45. 'DenseNet121':
  46. 'https://paddle-imagenet-models-name.bj.bcebos.com/DenseNet121_pretrained.tar',
  47. 'DenseNet161':
  48. 'https://paddle-imagenet-models-name.bj.bcebos.com/DenseNet161_pretrained.tar',
  49. 'DenseNet201':
  50. 'https://paddle-imagenet-models-name.bj.bcebos.com/DenseNet201_pretrained.tar',
  51. 'DetResNet50':
  52. 'https://paddle-imagenet-models-name.bj.bcebos.com/ResNet50_cos_pretrained.tar',
  53. 'SegXception41':
  54. 'https://paddle-imagenet-models-name.bj.bcebos.com/Xception41_deeplab_pretrained.tar',
  55. 'SegXception65':
  56. 'https://paddle-imagenet-models-name.bj.bcebos.com/Xception65_deeplab_pretrained.tar',
  57. 'ShuffleNetV2':
  58. 'https://paddle-imagenet-models-name.bj.bcebos.com/ShuffleNetV2_pretrained.tar',
  59. 'HRNet_W18':
  60. 'https://paddle-imagenet-models-name.bj.bcebos.com/HRNet_W18_C_pretrained.tar',
  61. 'HRNet_W30':
  62. 'https://paddle-imagenet-models-name.bj.bcebos.com/HRNet_W30_C_pretrained.tar',
  63. 'HRNet_W32':
  64. 'https://paddle-imagenet-models-name.bj.bcebos.com/HRNet_W32_C_pretrained.tar',
  65. 'HRNet_W40':
  66. 'https://paddle-imagenet-models-name.bj.bcebos.com/HRNet_W40_C_pretrained.tar',
  67. 'HRNet_W48':
  68. 'https://paddle-imagenet-models-name.bj.bcebos.com/HRNet_W48_C_pretrained.tar',
  69. 'HRNet_W60':
  70. 'https://paddle-imagenet-models-name.bj.bcebos.com/HRNet_W60_C_pretrained.tar',
  71. 'HRNet_W64':
  72. 'https://paddle-imagenet-models-name.bj.bcebos.com/HRNet_W64_C_pretrained.tar',
  73. 'AlexNet':
  74. 'http://paddle-imagenet-models-name.bj.bcebos.com/AlexNet_pretrained.tar'
  75. }
  76. coco_pretrain = {
  77. 'YOLOv3_DarkNet53_COCO':
  78. 'https://paddlemodels.bj.bcebos.com/object_detection/yolov3_darknet.tar',
  79. 'YOLOv3_MobileNetV1_COCO':
  80. 'https://paddlemodels.bj.bcebos.com/object_detection/yolov3_mobilenet_v1.tar',
  81. 'YOLOv3_MobileNetV3_large_COCO':
  82. 'https://paddlemodels.bj.bcebos.com/object_detection/yolov3_mobilenet_v3.pdparams',
  83. 'YOLOv3_ResNet34_COCO':
  84. 'https://paddlemodels.bj.bcebos.com/object_detection/yolov3_r34.tar',
  85. 'YOLOv3_ResNet50_vd_COCO':
  86. 'https://paddlemodels.bj.bcebos.com/object_detection/yolov3_r50vd_dcn.tar',
  87. 'FasterRCNN_ResNet50_COCO':
  88. 'https://paddlemodels.bj.bcebos.com/object_detection/faster_rcnn_r50_fpn_2x.tar',
  89. 'FasterRCNN_ResNet50_vd_COCO':
  90. 'https://paddlemodels.bj.bcebos.com/object_detection/faster_rcnn_r50_vd_fpn_2x.tar',
  91. 'FasterRCNN_ResNet101_COCO':
  92. 'https://paddlemodels.bj.bcebos.com/object_detection/faster_rcnn_r101_fpn_2x.tar',
  93. 'FasterRCNN_ResNet101_vd_COCO':
  94. 'https://paddlemodels.bj.bcebos.com/object_detection/faster_rcnn_r101_vd_fpn_2x.tar',
  95. 'FasterRCNN_HRNet_W18_COCO':
  96. 'https://paddlemodels.bj.bcebos.com/object_detection/faster_rcnn_hrnetv2p_w18_2x.tar',
  97. 'MaskRCNN_ResNet50_COCO':
  98. 'https://paddlemodels.bj.bcebos.com/object_detection/mask_rcnn_r50_fpn_2x.tar',
  99. 'MaskRCNN_ResNet50_vd_COCO':
  100. 'https://paddlemodels.bj.bcebos.com/object_detection/mask_rcnn_r50_vd_fpn_2x.tar',
  101. 'MaskRCNN_ResNet101_COCO':
  102. 'https://paddlemodels.bj.bcebos.com/object_detection/mask_rcnn_r101_fpn_1x.tar',
  103. 'MaskRCNN_ResNet101_vd_COCO':
  104. 'https://paddlemodels.bj.bcebos.com/object_detection/mask_rcnn_r101_vd_fpn_1x.tar',
  105. 'UNet_COCO': 'https://paddleseg.bj.bcebos.com/models/unet_coco_v3.tgz',
  106. 'DeepLabv3p_MobileNetV2_x1.0_COCO':
  107. 'https://bj.bcebos.com/v1/paddleseg/deeplab_mobilenet_x1_0_coco.tgz',
  108. 'DeepLabv3p_Xception65_COCO':
  109. 'https://paddleseg.bj.bcebos.com/models/xception65_coco.tgz'
  110. }
  111. cityscapes_pretrain = {
  112. 'DeepLabv3p_MobileNetV2_x1.0_CITYSCAPES':
  113. 'https://paddleseg.bj.bcebos.com/models/mobilenet_cityscapes.tgz',
  114. 'DeepLabv3p_Xception65_CITYSCAPES':
  115. 'https://paddleseg.bj.bcebos.com/models/xception65_bn_cityscapes.tgz',
  116. 'HRNet_W18_CITYSCAPES':
  117. 'https://paddleseg.bj.bcebos.com/models/hrnet_w18_bn_cityscapes.tgz'
  118. }
  119. def get_pretrain_weights(flag, class_name, backbone, save_dir):
  120. if flag is None:
  121. return None
  122. elif osp.isdir(flag):
  123. return flag
  124. elif osp.isfile(flag):
  125. return flag
  126. warning_info = "{} does not support to be finetuned with weights pretrained on the {} dataset, so pretrain_weights is forced to be set to {}"
  127. if flag == 'COCO':
  128. if class_name == "FasterRCNN" and backbone in ['ResNet18'] or \
  129. class_name == "MaskRCNN" and backbone in ['ResNet18', 'HRNet_W18'] or \
  130. class_name == 'DeepLabv3p' and backbone in ['Xception41', 'MobileNetV2_x0.25', 'MobileNetV2_x0.5', 'MobileNetV2_x1.5', 'MobileNetV2_x2.0']:
  131. model_name = '{}_{}'.format(class_name, backbone)
  132. logging.warning(warning_info.format(model_name, flag, 'IMAGENET'))
  133. flag = 'IMAGENET'
  134. elif class_name == 'HRNet':
  135. logging.warning(warning_info.format(class_name, flag, 'IMAGENET'))
  136. flag = 'IMAGENET'
  137. elif flag == 'CITYSCAPES':
  138. model_name = '{}_{}'.format(class_name, backbone)
  139. if class_name == 'UNet':
  140. logging.warning(warning_info.format(class_name, flag, 'COCO'))
  141. flag = 'COCO'
  142. if class_name == 'HRNet' and backbone.split('_')[
  143. -1] in ['W30', 'W32', 'W40', 'W48', 'W60', 'W64']:
  144. logging.warning(warning_info.format(backbone, flag, 'IMAGENET'))
  145. flag = 'IMAGENET'
  146. if class_name == 'DeepLabv3p' and backbone in [
  147. 'Xception41', 'MobileNetV2_x0.25', 'MobileNetV2_x0.5',
  148. 'MobileNetV2_x1.5', 'MobileNetV2_x2.0'
  149. ]:
  150. model_name = '{}_{}'.format(class_name, backbone)
  151. logging.warning(warning_info.format(model_name, flag, 'IMAGENET'))
  152. flag = 'IMAGENET'
  153. elif flag == 'IMAGENET' and class_name == 'UNet':
  154. logging.warning(warning_info.format(class_name, flag, 'COCO'))
  155. flag = 'COCO'
  156. if flag == 'IMAGENET':
  157. new_save_dir = save_dir
  158. if hasattr(paddlex, 'pretrain_dir'):
  159. new_save_dir = paddlex.pretrain_dir
  160. if backbone.startswith('Xception'):
  161. backbone = 'Seg{}'.format(backbone)
  162. elif backbone == 'MobileNetV2':
  163. backbone = 'MobileNetV2_x1.0'
  164. elif backbone == 'MobileNetV3_small_ssld':
  165. backbone = 'MobileNetV3_small_x1_0_ssld'
  166. elif backbone == 'MobileNetV3_large_ssld':
  167. backbone = 'MobileNetV3_large_x1_0_ssld'
  168. if class_name in ['YOLOv3', 'FasterRCNN', 'MaskRCNN']:
  169. if backbone == 'ResNet50':
  170. backbone = 'DetResNet50'
  171. assert backbone in image_pretrain, "There is not ImageNet pretrain weights for {}, you may try COCO.".format(
  172. backbone)
  173. # if backbone == 'AlexNet':
  174. # url = image_pretrain[backbone]
  175. # fname = osp.split(url)[-1].split('.')[0]
  176. # paddlex.utils.download_and_decompress(url, path=new_save_dir)
  177. # return osp.join(new_save_dir, fname)
  178. try:
  179. hub.download(backbone, save_path=new_save_dir)
  180. except Exception as e:
  181. if isinstance(e, hub.ResourceNotFoundError):
  182. raise Exception("Resource for backbone {} not found".format(
  183. backbone))
  184. elif isinstance(e, hub.ServerConnectionError):
  185. raise Exception(
  186. "Cannot get reource for backbone {}, please check your internet connecgtion"
  187. .format(backbone))
  188. else:
  189. raise Exception(
  190. "Unexpected error, please make sure paddlehub >= 1.6.2")
  191. return osp.join(new_save_dir, backbone)
  192. elif flag in ['COCO', 'CITYSCAPES']:
  193. new_save_dir = save_dir
  194. if hasattr(paddlex, 'pretrain_dir'):
  195. new_save_dir = paddlex.pretrain_dir
  196. if class_name in ['YOLOv3', 'FasterRCNN', 'MaskRCNN', 'DeepLabv3p']:
  197. backbone = '{}_{}'.format(class_name, backbone)
  198. backbone = "{}_{}".format(backbone, flag)
  199. if flag == 'COCO':
  200. url = coco_pretrain[backbone]
  201. elif flag == 'CITYSCAPES':
  202. url = cityscapes_pretrain[backbone]
  203. fname = osp.split(url)[-1].split('.')[0]
  204. # paddlex.utils.download_and_decompress(url, path=new_save_dir)
  205. # return osp.join(new_save_dir, fname)
  206. try:
  207. hub.download(backbone, save_path=new_save_dir)
  208. except Exception as e:
  209. if isinstance(hub.ResourceNotFoundError):
  210. raise Exception("Resource for backbone {} not found".format(
  211. backbone))
  212. elif isinstance(hub.ServerConnectionError):
  213. raise Exception(
  214. "Cannot get reource for backbone {}, please check your internet connecgtion"
  215. .format(backbone))
  216. else:
  217. raise Exception(
  218. "Unexpected error, please make sure paddlehub >= 1.6.2")
  219. return osp.join(new_save_dir, backbone)
  220. else:
  221. raise Exception(
  222. "pretrain_weights need to be defined as directory path or 'IMAGENET' or 'COCO' or 'Cityscapes' (download pretrain weights automatically)."
  223. )