unet.py 7.6 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153
  1. #copyright (c) 2020 PaddlePaddle Authors. All Rights Reserve.
  2. #
  3. #Licensed under the Apache License, Version 2.0 (the "License");
  4. #you may not use this file except in compliance with the License.
  5. #You may obtain a copy of the License at
  6. #
  7. # http://www.apache.org/licenses/LICENSE-2.0
  8. #
  9. #Unless required by applicable law or agreed to in writing, software
  10. #distributed under the License is distributed on an "AS IS" BASIS,
  11. #WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  12. #See the License for the specific language governing permissions and
  13. #limitations under the License.
  14. from __future__ import absolute_import
  15. import paddlex
  16. from collections import OrderedDict
  17. from .deeplabv3p import DeepLabv3p
  18. class UNet(DeepLabv3p):
  19. """实现UNet网络的构建并进行训练、评估、预测和模型导出。
  20. Args:
  21. num_classes (int): 类别数。
  22. upsample_mode (str): UNet decode时采用的上采样方式,取值为'bilinear'时利用双线行差值进行上菜样,
  23. 当输入其他选项时则利用反卷积进行上菜样,默认为'bilinear'。
  24. use_bce_loss (bool): 是否使用bce loss作为网络的损失函数,只能用于两类分割。可与dice loss同时使用。默认False。
  25. use_dice_loss (bool): 是否使用dice loss作为网络的损失函数,只能用于两类分割,可与bce loss同时使用。
  26. 当use_bce_loss和use_dice_loss都为False时,使用交叉熵损失函数。默认False。
  27. class_weight (list/str): 交叉熵损失函数各类损失的权重。当class_weight为list的时候,长度应为
  28. num_classes。当class_weight为str时, weight.lower()应为'dynamic',这时会根据每一轮各类像素的比重
  29. 自行计算相应的权重,每一类的权重为:每类的比例 * num_classes。class_weight取默认值None是,各类的权重1,
  30. 即平时使用的交叉熵损失函数。
  31. ignore_index (int): label上忽略的值,label为ignore_index的像素不参与损失函数的计算。默认255。
  32. fixed_input_shape (list): 长度为2,维度为1的list,如:[640,720],用来固定模型输入:'image'的shape,默认为None。
  33. Raises:
  34. ValueError: use_bce_loss或use_dice_loss为真且num_calsses > 2。
  35. ValueError: class_weight为list, 但长度不等于num_class。
  36. class_weight为str, 但class_weight.low()不等于dynamic。
  37. TypeError: class_weight不为None时,其类型不是list或str。
  38. """
  39. def __init__(self,
  40. num_classes=2,
  41. upsample_mode='bilinear',
  42. use_bce_loss=False,
  43. use_dice_loss=False,
  44. class_weight=None,
  45. ignore_index=255,
  46. fixed_input_shape=None):
  47. self.init_params = locals()
  48. super(DeepLabv3p, self).__init__('segmenter')
  49. # dice_loss或bce_loss只适用两类分割中
  50. if num_classes > 2 and (use_bce_loss or use_dice_loss):
  51. raise ValueError(
  52. "dice loss and bce loss is only applicable to binary classfication"
  53. )
  54. if class_weight is not None:
  55. if isinstance(class_weight, list):
  56. if len(class_weight) != num_classes:
  57. raise ValueError(
  58. "Length of class_weight should be equal to number of classes"
  59. )
  60. elif isinstance(class_weight, str):
  61. if class_weight.lower() != 'dynamic':
  62. raise ValueError(
  63. "if class_weight is string, must be dynamic!")
  64. else:
  65. raise TypeError(
  66. 'Expect class_weight is a list or string but receive {}'.
  67. format(type(class_weight)))
  68. self.num_classes = num_classes
  69. self.upsample_mode = upsample_mode
  70. self.use_bce_loss = use_bce_loss
  71. self.use_dice_loss = use_dice_loss
  72. self.class_weight = class_weight
  73. self.ignore_index = ignore_index
  74. self.labels = None
  75. self.fixed_input_shape = fixed_input_shape
  76. def build_net(self, mode='train'):
  77. model = paddlex.cv.nets.segmentation.UNet(
  78. self.num_classes,
  79. mode=mode,
  80. upsample_mode=self.upsample_mode,
  81. use_bce_loss=self.use_bce_loss,
  82. use_dice_loss=self.use_dice_loss,
  83. class_weight=self.class_weight,
  84. ignore_index=self.ignore_index,
  85. fixed_input_shape = self.fixed_input_shape)
  86. inputs = model.generate_inputs()
  87. model_out = model.build_net(inputs)
  88. outputs = OrderedDict()
  89. if mode == 'train':
  90. self.optimizer.minimize(model_out)
  91. outputs['loss'] = model_out
  92. elif mode == 'eval':
  93. outputs['loss'] = model_out[0]
  94. outputs['pred'] = model_out[1]
  95. outputs['label'] = model_out[2]
  96. outputs['mask'] = model_out[3]
  97. else:
  98. outputs['pred'] = model_out[0]
  99. outputs['logit'] = model_out[1]
  100. return inputs, outputs
  101. def train(self,
  102. num_epochs,
  103. train_dataset,
  104. train_batch_size=2,
  105. eval_dataset=None,
  106. save_interval_epochs=1,
  107. log_interval_steps=2,
  108. save_dir='output',
  109. pretrain_weights='COCO',
  110. optimizer=None,
  111. learning_rate=0.01,
  112. lr_decay_power=0.9,
  113. use_vdl=False,
  114. sensitivities_file=None,
  115. eval_metric_loss=0.05):
  116. """训练。
  117. Args:
  118. num_epochs (int): 训练迭代轮数。
  119. train_dataset (paddlex.datasets): 训练数据读取器。
  120. train_batch_size (int): 训练数据batch大小。同时作为验证数据batch大小。默认2。
  121. eval_dataset (paddlex.datasets): 评估数据读取器。
  122. save_interval_epochs (int): 模型保存间隔(单位:迭代轮数)。默认为1。
  123. log_interval_steps (int): 训练日志输出间隔(单位:迭代次数)。默认为2。
  124. save_dir (str): 模型保存路径。默认'output'。
  125. pretrain_weights (str): 若指定为路径时,则加载路径下预训练模型;若为字符串'COCO',
  126. 则自动下载在COCO图片数据上预训练的模型权重;若为None,则不使用预训练模型。默认为'COCO'。
  127. optimizer (paddle.fluid.optimizer): 优化器。当改参数为None时,使用默认的优化器:使用
  128. fluid.optimizer.Momentum优化方法,polynomial的学习率衰减策略。
  129. learning_rate (float): 默认优化器的初始学习率。默认0.01。
  130. lr_decay_power (float): 默认优化器学习率多项式衰减系数。默认0.9。
  131. use_vdl (bool): 是否使用VisualDL进行可视化。默认False。
  132. sensitivities_file (str): 若指定为路径时,则加载路径下敏感度信息进行裁剪;若为字符串'DEFAULT',
  133. 则自动下载在ImageNet图片数据上获得的敏感度信息进行裁剪;若为None,则不进行裁剪。默认为None。
  134. eval_metric_loss (float): 可容忍的精度损失。默认为0.05。
  135. Raises:
  136. ValueError: 模型从inference model进行加载。
  137. """
  138. return super(UNet, self).train(
  139. num_epochs, train_dataset, train_batch_size, eval_dataset,
  140. save_interval_epochs, log_interval_steps, save_dir,
  141. pretrain_weights, optimizer, learning_rate, lr_decay_power,
  142. use_vdl, sensitivities_file, eval_metric_loss)