model.py 10 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283
  1. #!/usr/bin/python
  2. # -*- encoding: utf-8 -*-
  3. import torch
  4. import torch.nn as nn
  5. import torch.nn.functional as F
  6. import torchvision
  7. from parsing_model.resnet import Resnet18
  8. # from modules.bn import InPlaceABNSync as BatchNorm2d
  9. class ConvBNReLU(nn.Module):
  10. def __init__(self, in_chan, out_chan, ks=3, stride=1, padding=1, *args, **kwargs):
  11. super(ConvBNReLU, self).__init__()
  12. self.conv = nn.Conv2d(in_chan,
  13. out_chan,
  14. kernel_size = ks,
  15. stride = stride,
  16. padding = padding,
  17. bias = False)
  18. self.bn = nn.BatchNorm2d(out_chan)
  19. self.init_weight()
  20. def forward(self, x):
  21. x = self.conv(x)
  22. x = F.relu(self.bn(x))
  23. return x
  24. def init_weight(self):
  25. for ly in self.children():
  26. if isinstance(ly, nn.Conv2d):
  27. nn.init.kaiming_normal_(ly.weight, a=1)
  28. if not ly.bias is None: nn.init.constant_(ly.bias, 0)
  29. class BiSeNetOutput(nn.Module):
  30. def __init__(self, in_chan, mid_chan, n_classes, *args, **kwargs):
  31. super(BiSeNetOutput, self).__init__()
  32. self.conv = ConvBNReLU(in_chan, mid_chan, ks=3, stride=1, padding=1)
  33. self.conv_out = nn.Conv2d(mid_chan, n_classes, kernel_size=1, bias=False)
  34. self.init_weight()
  35. def forward(self, x):
  36. x = self.conv(x)
  37. x = self.conv_out(x)
  38. return x
  39. def init_weight(self):
  40. for ly in self.children():
  41. if isinstance(ly, nn.Conv2d):
  42. nn.init.kaiming_normal_(ly.weight, a=1)
  43. if not ly.bias is None: nn.init.constant_(ly.bias, 0)
  44. def get_params(self):
  45. wd_params, nowd_params = [], []
  46. for name, module in self.named_modules():
  47. if isinstance(module, nn.Linear) or isinstance(module, nn.Conv2d):
  48. wd_params.append(module.weight)
  49. if not module.bias is None:
  50. nowd_params.append(module.bias)
  51. elif isinstance(module, nn.BatchNorm2d):
  52. nowd_params += list(module.parameters())
  53. return wd_params, nowd_params
  54. class AttentionRefinementModule(nn.Module):
  55. def __init__(self, in_chan, out_chan, *args, **kwargs):
  56. super(AttentionRefinementModule, self).__init__()
  57. self.conv = ConvBNReLU(in_chan, out_chan, ks=3, stride=1, padding=1)
  58. self.conv_atten = nn.Conv2d(out_chan, out_chan, kernel_size= 1, bias=False)
  59. self.bn_atten = nn.BatchNorm2d(out_chan)
  60. self.sigmoid_atten = nn.Sigmoid()
  61. self.init_weight()
  62. def forward(self, x):
  63. feat = self.conv(x)
  64. atten = F.avg_pool2d(feat, feat.size()[2:])
  65. atten = self.conv_atten(atten)
  66. atten = self.bn_atten(atten)
  67. atten = self.sigmoid_atten(atten)
  68. out = torch.mul(feat, atten)
  69. return out
  70. def init_weight(self):
  71. for ly in self.children():
  72. if isinstance(ly, nn.Conv2d):
  73. nn.init.kaiming_normal_(ly.weight, a=1)
  74. if not ly.bias is None: nn.init.constant_(ly.bias, 0)
  75. class ContextPath(nn.Module):
  76. def __init__(self, *args, **kwargs):
  77. super(ContextPath, self).__init__()
  78. self.resnet = Resnet18()
  79. self.arm16 = AttentionRefinementModule(256, 128)
  80. self.arm32 = AttentionRefinementModule(512, 128)
  81. self.conv_head32 = ConvBNReLU(128, 128, ks=3, stride=1, padding=1)
  82. self.conv_head16 = ConvBNReLU(128, 128, ks=3, stride=1, padding=1)
  83. self.conv_avg = ConvBNReLU(512, 128, ks=1, stride=1, padding=0)
  84. self.init_weight()
  85. def forward(self, x):
  86. H0, W0 = x.size()[2:]
  87. feat8, feat16, feat32 = self.resnet(x)
  88. H8, W8 = feat8.size()[2:]
  89. H16, W16 = feat16.size()[2:]
  90. H32, W32 = feat32.size()[2:]
  91. avg = F.avg_pool2d(feat32, feat32.size()[2:])
  92. avg = self.conv_avg(avg)
  93. avg_up = F.interpolate(avg, (H32, W32), mode='nearest')
  94. feat32_arm = self.arm32(feat32)
  95. feat32_sum = feat32_arm + avg_up
  96. feat32_up = F.interpolate(feat32_sum, (H16, W16), mode='nearest')
  97. feat32_up = self.conv_head32(feat32_up)
  98. feat16_arm = self.arm16(feat16)
  99. feat16_sum = feat16_arm + feat32_up
  100. feat16_up = F.interpolate(feat16_sum, (H8, W8), mode='nearest')
  101. feat16_up = self.conv_head16(feat16_up)
  102. return feat8, feat16_up, feat32_up # x8, x8, x16
  103. def init_weight(self):
  104. for ly in self.children():
  105. if isinstance(ly, nn.Conv2d):
  106. nn.init.kaiming_normal_(ly.weight, a=1)
  107. if not ly.bias is None: nn.init.constant_(ly.bias, 0)
  108. def get_params(self):
  109. wd_params, nowd_params = [], []
  110. for name, module in self.named_modules():
  111. if isinstance(module, (nn.Linear, nn.Conv2d)):
  112. wd_params.append(module.weight)
  113. if not module.bias is None:
  114. nowd_params.append(module.bias)
  115. elif isinstance(module, nn.BatchNorm2d):
  116. nowd_params += list(module.parameters())
  117. return wd_params, nowd_params
  118. ### This is not used, since I replace this with the resnet feature with the same size
  119. class SpatialPath(nn.Module):
  120. def __init__(self, *args, **kwargs):
  121. super(SpatialPath, self).__init__()
  122. self.conv1 = ConvBNReLU(3, 64, ks=7, stride=2, padding=3)
  123. self.conv2 = ConvBNReLU(64, 64, ks=3, stride=2, padding=1)
  124. self.conv3 = ConvBNReLU(64, 64, ks=3, stride=2, padding=1)
  125. self.conv_out = ConvBNReLU(64, 128, ks=1, stride=1, padding=0)
  126. self.init_weight()
  127. def forward(self, x):
  128. feat = self.conv1(x)
  129. feat = self.conv2(feat)
  130. feat = self.conv3(feat)
  131. feat = self.conv_out(feat)
  132. return feat
  133. def init_weight(self):
  134. for ly in self.children():
  135. if isinstance(ly, nn.Conv2d):
  136. nn.init.kaiming_normal_(ly.weight, a=1)
  137. if not ly.bias is None: nn.init.constant_(ly.bias, 0)
  138. def get_params(self):
  139. wd_params, nowd_params = [], []
  140. for name, module in self.named_modules():
  141. if isinstance(module, nn.Linear) or isinstance(module, nn.Conv2d):
  142. wd_params.append(module.weight)
  143. if not module.bias is None:
  144. nowd_params.append(module.bias)
  145. elif isinstance(module, nn.BatchNorm2d):
  146. nowd_params += list(module.parameters())
  147. return wd_params, nowd_params
  148. class FeatureFusionModule(nn.Module):
  149. def __init__(self, in_chan, out_chan, *args, **kwargs):
  150. super(FeatureFusionModule, self).__init__()
  151. self.convblk = ConvBNReLU(in_chan, out_chan, ks=1, stride=1, padding=0)
  152. self.conv1 = nn.Conv2d(out_chan,
  153. out_chan//4,
  154. kernel_size = 1,
  155. stride = 1,
  156. padding = 0,
  157. bias = False)
  158. self.conv2 = nn.Conv2d(out_chan//4,
  159. out_chan,
  160. kernel_size = 1,
  161. stride = 1,
  162. padding = 0,
  163. bias = False)
  164. self.relu = nn.ReLU(inplace=True)
  165. self.sigmoid = nn.Sigmoid()
  166. self.init_weight()
  167. def forward(self, fsp, fcp):
  168. fcat = torch.cat([fsp, fcp], dim=1)
  169. feat = self.convblk(fcat)
  170. atten = F.avg_pool2d(feat, feat.size()[2:])
  171. atten = self.conv1(atten)
  172. atten = self.relu(atten)
  173. atten = self.conv2(atten)
  174. atten = self.sigmoid(atten)
  175. feat_atten = torch.mul(feat, atten)
  176. feat_out = feat_atten + feat
  177. return feat_out
  178. def init_weight(self):
  179. for ly in self.children():
  180. if isinstance(ly, nn.Conv2d):
  181. nn.init.kaiming_normal_(ly.weight, a=1)
  182. if not ly.bias is None: nn.init.constant_(ly.bias, 0)
  183. def get_params(self):
  184. wd_params, nowd_params = [], []
  185. for name, module in self.named_modules():
  186. if isinstance(module, nn.Linear) or isinstance(module, nn.Conv2d):
  187. wd_params.append(module.weight)
  188. if not module.bias is None:
  189. nowd_params.append(module.bias)
  190. elif isinstance(module, nn.BatchNorm2d):
  191. nowd_params += list(module.parameters())
  192. return wd_params, nowd_params
  193. class BiSeNet(nn.Module):
  194. def __init__(self, n_classes, *args, **kwargs):
  195. super(BiSeNet, self).__init__()
  196. self.cp = ContextPath()
  197. ## here self.sp is deleted
  198. self.ffm = FeatureFusionModule(256, 256)
  199. self.conv_out = BiSeNetOutput(256, 256, n_classes)
  200. self.conv_out16 = BiSeNetOutput(128, 64, n_classes)
  201. self.conv_out32 = BiSeNetOutput(128, 64, n_classes)
  202. self.init_weight()
  203. def forward(self, x):
  204. H, W = x.size()[2:]
  205. feat_res8, feat_cp8, feat_cp16 = self.cp(x) # here return res3b1 feature
  206. feat_sp = feat_res8 # use res3b1 feature to replace spatial path feature
  207. feat_fuse = self.ffm(feat_sp, feat_cp8)
  208. feat_out = self.conv_out(feat_fuse)
  209. feat_out16 = self.conv_out16(feat_cp8)
  210. feat_out32 = self.conv_out32(feat_cp16)
  211. feat_out = F.interpolate(feat_out, (H, W), mode='bilinear', align_corners=True)
  212. feat_out16 = F.interpolate(feat_out16, (H, W), mode='bilinear', align_corners=True)
  213. feat_out32 = F.interpolate(feat_out32, (H, W), mode='bilinear', align_corners=True)
  214. return feat_out, feat_out16, feat_out32
  215. def init_weight(self):
  216. for ly in self.children():
  217. if isinstance(ly, nn.Conv2d):
  218. nn.init.kaiming_normal_(ly.weight, a=1)
  219. if not ly.bias is None: nn.init.constant_(ly.bias, 0)
  220. def get_params(self):
  221. wd_params, nowd_params, lr_mul_wd_params, lr_mul_nowd_params = [], [], [], []
  222. for name, child in self.named_children():
  223. child_wd_params, child_nowd_params = child.get_params()
  224. if isinstance(child, FeatureFusionModule) or isinstance(child, BiSeNetOutput):
  225. lr_mul_wd_params += child_wd_params
  226. lr_mul_nowd_params += child_nowd_params
  227. else:
  228. wd_params += child_wd_params
  229. nowd_params += child_nowd_params
  230. return wd_params, nowd_params, lr_mul_wd_params, lr_mul_nowd_params
  231. if __name__ == "__main__":
  232. net = BiSeNet(19)
  233. net.cuda()
  234. net.eval()
  235. in_ten = torch.randn(16, 3, 640, 480).cuda()
  236. out, out16, out32 = net(in_ten)
  237. print(out.shape)
  238. net.get_params()