fcn.py 4.8 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135
  1. # Copyright (c) 2020 PaddlePaddle Authors. All Rights Reserved.
  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. import paddle.nn as nn
  15. import paddle.nn.functional as F
  16. import paddle
  17. from paddlex.paddleseg import utils
  18. from paddlex.paddleseg.cvlibs import manager, param_init
  19. from paddlex.paddleseg.models import layers
  20. @manager.MODELS.add_component
  21. class FCN(nn.Layer):
  22. """
  23. A simple implementation for FCN based on PaddlePaddle.
  24. The original article refers to
  25. Evan Shelhamer, et, al. "Fully Convolutional Networks for Semantic Segmentation"
  26. (https://arxiv.org/abs/1411.4038).
  27. Args:
  28. num_classes (int): The unique number of target classes.
  29. backbone (paddle.nn.Layer): Backbone networks.
  30. backbone_indices (tuple, optional): The values in the tuple indicate the indices of output of backbone.
  31. Default: (-1, ).
  32. channels (int, optional): The channels between conv layer and the last layer of FCNHead.
  33. If None, it will be the number of channels of input features. Default: None.
  34. align_corners (bool): An argument of F.interpolate. It should be set to False when the output size of feature
  35. is even, e.g. 1024x512, otherwise it is True, e.g. 769x769. Default: False.
  36. pretrained (str, optional): The path or url of pretrained model. Default: None
  37. """
  38. def __init__(self,
  39. num_classes,
  40. backbone,
  41. backbone_indices=(-1, ),
  42. channels=None,
  43. align_corners=False,
  44. pretrained=None):
  45. super(FCN, self).__init__()
  46. self.backbone = backbone
  47. backbone_channels = [
  48. backbone.feat_channels[i] for i in backbone_indices
  49. ]
  50. self.head = FCNHead(num_classes, backbone_indices, backbone_channels,
  51. channels)
  52. self.align_corners = align_corners
  53. self.pretrained = pretrained
  54. self.init_weight()
  55. def forward(self, x):
  56. feat_list = self.backbone(x)
  57. logit_list = self.head(feat_list)
  58. return [
  59. F.interpolate(
  60. logit,
  61. paddle.shape(x)[2:],
  62. mode='bilinear',
  63. align_corners=self.align_corners) for logit in logit_list
  64. ]
  65. def init_weight(self):
  66. if self.pretrained is not None:
  67. utils.load_entire_model(self, self.pretrained)
  68. class FCNHead(nn.Layer):
  69. """
  70. A simple implementation for FCNHead based on PaddlePaddle
  71. Args:
  72. num_classes (int): The unique number of target classes.
  73. backbone_indices (tuple, optional): The values in the tuple indicate the indices of output of backbone.
  74. Default: (-1, ).
  75. channels (int, optional): The channels between conv layer and the last layer of FCNHead.
  76. If None, it will be the number of channels of input features. Default: None.
  77. pretrained (str, optional): The path of pretrained model. Default: None
  78. """
  79. def __init__(self,
  80. num_classes,
  81. backbone_indices=(-1, ),
  82. backbone_channels=(270, ),
  83. channels=None):
  84. super(FCNHead, self).__init__()
  85. self.num_classes = num_classes
  86. self.backbone_indices = backbone_indices
  87. if channels is None:
  88. channels = backbone_channels[0]
  89. self.conv_1 = layers.ConvBNReLU(
  90. in_channels=backbone_channels[0],
  91. out_channels=channels,
  92. kernel_size=1,
  93. padding='same',
  94. stride=1)
  95. self.cls = nn.Conv2D(
  96. in_channels=channels,
  97. out_channels=self.num_classes,
  98. kernel_size=1,
  99. stride=1,
  100. padding=0)
  101. self.init_weight()
  102. def forward(self, feat_list):
  103. logit_list = []
  104. x = feat_list[self.backbone_indices[0]]
  105. x = self.conv_1(x)
  106. logit = self.cls(x)
  107. logit_list.append(logit)
  108. return logit_list
  109. def init_weight(self):
  110. for layer in self.sublayers():
  111. if isinstance(layer, nn.Conv2D):
  112. param_init.normal_init(layer.weight, std=0.001)
  113. elif isinstance(layer, (nn.BatchNorm, nn.SyncBatchNorm)):
  114. param_init.constant_init(layer.weight, value=1.0)
  115. param_init.constant_init(layer.bias, value=0.0)