custom_model.py 7.1 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227
  1. import paddle
  2. import paddle.nn as nn
  3. import paddle.nn.functional as F
  4. import paddlers
  5. from paddlers.rs_models.cd.layers import Conv3x3, MaxPool2x2, ConvTransposed3x3, Identity
  6. from paddlers.rs_models.cd.layers import ChannelAttention
  7. from attach_tools import Attach
  8. attach = Attach.to(paddlers.rs_models.cd)
  9. @attach
  10. class CustomModel(nn.Layer):
  11. def __init__(self,
  12. in_channels,
  13. num_classes,
  14. att_types='ct',
  15. use_dropout=False):
  16. super(CustomModel, self).__init__()
  17. C1, C2, C3, C4, C5 = 16, 32, 64, 128, 256
  18. self.use_dropout = use_dropout
  19. self.conv11 = Conv3x3(in_channels, C1, norm=True, act=True)
  20. self.do11 = self._make_dropout()
  21. self.conv12 = Conv3x3(C1, C1, norm=True, act=True)
  22. self.do12 = self._make_dropout()
  23. self.pool1 = MaxPool2x2()
  24. self.conv21 = Conv3x3(C1, C2, norm=True, act=True)
  25. self.do21 = self._make_dropout()
  26. self.conv22 = Conv3x3(C2, C2, norm=True, act=True)
  27. self.do22 = self._make_dropout()
  28. self.pool2 = MaxPool2x2()
  29. self.conv31 = Conv3x3(C2, C3, norm=True, act=True)
  30. self.do31 = self._make_dropout()
  31. self.conv32 = Conv3x3(C3, C3, norm=True, act=True)
  32. self.do32 = self._make_dropout()
  33. self.conv33 = Conv3x3(C3, C3, norm=True, act=True)
  34. self.do33 = self._make_dropout()
  35. self.pool3 = MaxPool2x2()
  36. self.conv41 = Conv3x3(C3, C4, norm=True, act=True)
  37. self.do41 = self._make_dropout()
  38. self.conv42 = Conv3x3(C4, C4, norm=True, act=True)
  39. self.do42 = self._make_dropout()
  40. self.conv43 = Conv3x3(C4, C4, norm=True, act=True)
  41. self.do43 = self._make_dropout()
  42. self.pool4 = MaxPool2x2()
  43. self.upconv4 = ConvTransposed3x3(C4, C4, output_padding=1)
  44. self.conv43d = Conv3x3(C5 + C4, C4, norm=True, act=True)
  45. self.do43d = self._make_dropout()
  46. self.conv42d = Conv3x3(C4, C4, norm=True, act=True)
  47. self.do42d = self._make_dropout()
  48. self.conv41d = Conv3x3(C4, C3, norm=True, act=True)
  49. self.do41d = self._make_dropout()
  50. self.upconv3 = ConvTransposed3x3(C3, C3, output_padding=1)
  51. self.conv33d = Conv3x3(C4 + C3, C3, norm=True, act=True)
  52. self.do33d = self._make_dropout()
  53. self.conv32d = Conv3x3(C3, C3, norm=True, act=True)
  54. self.do32d = self._make_dropout()
  55. self.conv31d = Conv3x3(C3, C2, norm=True, act=True)
  56. self.do31d = self._make_dropout()
  57. self.upconv2 = ConvTransposed3x3(C2, C2, output_padding=1)
  58. self.conv22d = Conv3x3(C3 + C2, C2, norm=True, act=True)
  59. self.do22d = self._make_dropout()
  60. self.conv21d = Conv3x3(C2, C1, norm=True, act=True)
  61. self.do21d = self._make_dropout()
  62. self.upconv1 = ConvTransposed3x3(C1, C1, output_padding=1)
  63. self.conv12d = Conv3x3(C2 + C1, C1, norm=True, act=True)
  64. self.do12d = self._make_dropout()
  65. self.conv11d = Conv3x3(C1, num_classes)
  66. self.init_weight()
  67. self.att4 = MixedAttention(C4, att_types)
  68. def forward(self, t1, t2):
  69. # Encode t1
  70. # Stage 1
  71. x11 = self.do11(self.conv11(t1))
  72. x12_1 = self.do12(self.conv12(x11))
  73. x1p = self.pool1(x12_1)
  74. # Stage 2
  75. x21 = self.do21(self.conv21(x1p))
  76. x22_1 = self.do22(self.conv22(x21))
  77. x2p = self.pool2(x22_1)
  78. # Stage 3
  79. x31 = self.do31(self.conv31(x2p))
  80. x32 = self.do32(self.conv32(x31))
  81. x33_1 = self.do33(self.conv33(x32))
  82. x3p = self.pool3(x33_1)
  83. # Stage 4
  84. x41 = self.do41(self.conv41(x3p))
  85. x42 = self.do42(self.conv42(x41))
  86. x43_1 = self.do43(self.conv43(x42))
  87. x4p = self.pool4(x43_1)
  88. # Encode t2
  89. # Stage 1
  90. x11 = self.do11(self.conv11(t2))
  91. x12_2 = self.do12(self.conv12(x11))
  92. x1p = self.pool1(x12_2)
  93. # Stage 2
  94. x21 = self.do21(self.conv21(x1p))
  95. x22_2 = self.do22(self.conv22(x21))
  96. x2p = self.pool2(x22_2)
  97. # Stage 3
  98. x31 = self.do31(self.conv31(x2p))
  99. x32 = self.do32(self.conv32(x31))
  100. x33_2 = self.do33(self.conv33(x32))
  101. x3p = self.pool3(x33_2)
  102. # Stage 4
  103. x41 = self.do41(self.conv41(x3p))
  104. x42 = self.do42(self.conv42(x41))
  105. x43_2 = self.do43(self.conv43(x42))
  106. x4p = self.pool4(x43_2)
  107. # Decode
  108. # Stage 4d
  109. x4d = self.upconv4(x4p)
  110. pad4 = (0, x43_1.shape[3] - x4d.shape[3], 0,
  111. x43_1.shape[2] - x4d.shape[2])
  112. x4d = F.pad(x4d, pad=pad4, mode='replicate')
  113. x43_1, x43_2 = self.att4(x43_1, x43_2)
  114. x4d = paddle.concat([x4d, x43_1, x43_2], 1)
  115. x43d = self.do43d(self.conv43d(x4d))
  116. x42d = self.do42d(self.conv42d(x43d))
  117. x41d = self.do41d(self.conv41d(x42d))
  118. # Stage 3d
  119. x3d = self.upconv3(x41d)
  120. pad3 = (0, x33_1.shape[3] - x3d.shape[3], 0,
  121. x33_1.shape[2] - x3d.shape[2])
  122. x3d = F.pad(x3d, pad=pad3, mode='replicate')
  123. x3d = paddle.concat([x3d, x33_1, x33_2], 1)
  124. x33d = self.do33d(self.conv33d(x3d))
  125. x32d = self.do32d(self.conv32d(x33d))
  126. x31d = self.do31d(self.conv31d(x32d))
  127. # Stage 2d
  128. x2d = self.upconv2(x31d)
  129. pad2 = (0, x22_1.shape[3] - x2d.shape[3], 0,
  130. x22_1.shape[2] - x2d.shape[2])
  131. x2d = F.pad(x2d, pad=pad2, mode='replicate')
  132. x2d = paddle.concat([x2d, x22_1, x22_2], 1)
  133. x22d = self.do22d(self.conv22d(x2d))
  134. x21d = self.do21d(self.conv21d(x22d))
  135. # Stage 1d
  136. x1d = self.upconv1(x21d)
  137. pad1 = (0, x12_1.shape[3] - x1d.shape[3], 0,
  138. x12_1.shape[2] - x1d.shape[2])
  139. x1d = F.pad(x1d, pad=pad1, mode='replicate')
  140. x1d = paddle.concat([x1d, x12_1, x12_2], 1)
  141. x12d = self.do12d(self.conv12d(x1d))
  142. x11d = self.conv11d(x12d)
  143. return [x11d]
  144. def init_weight(self):
  145. pass
  146. def _make_dropout(self):
  147. if self.use_dropout:
  148. return nn.Dropout2D(p=0.2)
  149. else:
  150. return Identity()
  151. class MixedAttention(nn.Layer):
  152. def __init__(self, in_channels, att_types='ct'):
  153. super(MixedAttention, self).__init__()
  154. self.att_types = att_types
  155. if self.has_att_c:
  156. self.att_c = ChannelAttention(in_channels, ratio=1)
  157. else:
  158. self.att_c = Identity()
  159. if self.has_att_t:
  160. self.att_t = ChannelAttention(2, ratio=1)
  161. else:
  162. self.att_t = Identity()
  163. def forward(self, x1, x2):
  164. if self.has_att_c:
  165. x1 = (1 + self.att_c(x1)) * x1
  166. x2 = (1 + self.att_c(x2)) * x2
  167. if self.has_att_t:
  168. b, c = x1.shape[:2]
  169. y = paddle.stack([x1, x2], axis=2)
  170. y = paddle.flatten(y, stop_axis=1)
  171. y = (1 + self.att_t(y)) * y
  172. y = y.reshape((b, c, 2, *y.shape[2:]))
  173. y1, y2 = y[:, :, 0], y[:, :, 1]
  174. else:
  175. y1, y2 = x1, x2
  176. return y1, y2
  177. @property
  178. def has_att_c(self):
  179. return 'c' in self.att_types
  180. @property
  181. def has_att_t(self):
  182. return 't' in self.att_types