yolo_fpn.py 37 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555556557558559560561562563564565566567568569570571572573574575576577578579580581582583584585586587588589590591592593594595596597598599600601602603604605606607608609610611612613614615616617618619620621622623624625626627628629630631632633634635636637638639640641642643644645646647648649650651652653654655656657658659660661662663664665666667668669670671672673674675676677678679680681682683684685686687688689690691692693694695696697698699700701702703704705706707708709710711712713714715716717718719720721722723724725726727728729730731732733734735736737738739740741742743744745746747748749750751752753754755756757758759760761762763764765766767768769770771772773774775776777778779780781782783784785786787788789790791792793794795796797798799800801802803804805806807808809810811812813814815816817818819820821822823824825826827828829830831832833834835836837838839840841842843844845846847848849850851852853854855856857858859860861862863864865866867868869870871872873874875876877878879880881882883884885886887888889890891892893894895896897898899900901902903904905906907908909910911912913914915916917918919920921922923924925926927928929930931932933934935936937938939940941942943944945946947948949950951952953954955956957958959960961962963964965966967968969970971972973974975976977978979980981982983984985986987988989990991992993994995996997998999100010011002100310041005100610071008100910101011101210131014101510161017101810191020102110221023102410251026102710281029103010311032103310341035103610371038103910401041104210431044104510461047104810491050105110521053105410551056105710581059106010611062106310641065106610671068106910701071107210731074107510761077107810791080108110821083108410851086108710881089109010911092109310941095109610971098109911001101
  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
  15. import paddle.nn as nn
  16. import paddle.nn.functional as F
  17. from paddlex.ppdet.core.workspace import register, serializable
  18. from paddlex.ppdet.modeling.layers import DropBlock
  19. from paddlex.ppdet.modeling.ops import get_act_fn
  20. from ..backbones.darknet import ConvBNLayer
  21. from ..shape_spec import ShapeSpec
  22. from ..backbones.csp_darknet import BaseConv, DWConv, CSPLayer
  23. __all__ = [
  24. 'YOLOv3FPN', 'PPYOLOFPN', 'PPYOLOTinyFPN', 'PPYOLOPAN', 'YOLOCSPPAN'
  25. ]
  26. def add_coord(x, data_format):
  27. b = paddle.shape(x)[0]
  28. if data_format == 'NCHW':
  29. h, w = x.shape[2], x.shape[3]
  30. else:
  31. h, w = x.shape[1], x.shape[2]
  32. gx = paddle.cast(paddle.arange(w) / ((w - 1.) * 2.0) - 1., x.dtype)
  33. gy = paddle.cast(paddle.arange(h) / ((h - 1.) * 2.0) - 1., x.dtype)
  34. if data_format == 'NCHW':
  35. gx = gx.reshape([1, 1, 1, w]).expand([b, 1, h, w])
  36. gy = gy.reshape([1, 1, h, 1]).expand([b, 1, h, w])
  37. else:
  38. gx = gx.reshape([1, 1, w, 1]).expand([b, h, w, 1])
  39. gy = gy.reshape([1, h, 1, 1]).expand([b, h, w, 1])
  40. gx.stop_gradient = True
  41. gy.stop_gradient = True
  42. return gx, gy
  43. class YoloDetBlock(nn.Layer):
  44. def __init__(self,
  45. ch_in,
  46. channel,
  47. norm_type,
  48. freeze_norm=False,
  49. name='',
  50. data_format='NCHW'):
  51. """
  52. YOLODetBlock layer for yolov3, see https://arxiv.org/abs/1804.02767
  53. Args:
  54. ch_in (int): input channel
  55. channel (int): base channel
  56. norm_type (str): batch norm type
  57. freeze_norm (bool): whether to freeze norm, default False
  58. name (str): layer name
  59. data_format (str): data format, NCHW or NHWC
  60. """
  61. super(YoloDetBlock, self).__init__()
  62. self.ch_in = ch_in
  63. self.channel = channel
  64. assert channel % 2 == 0, \
  65. "channel {} cannot be divided by 2".format(channel)
  66. conv_def = [
  67. ['conv0', ch_in, channel, 1, '.0.0'],
  68. ['conv1', channel, channel * 2, 3, '.0.1'],
  69. ['conv2', channel * 2, channel, 1, '.1.0'],
  70. ['conv3', channel, channel * 2, 3, '.1.1'],
  71. ['route', channel * 2, channel, 1, '.2'],
  72. ]
  73. self.conv_module = nn.Sequential()
  74. for idx, (conv_name, ch_in, ch_out, filter_size,
  75. post_name) in enumerate(conv_def):
  76. self.conv_module.add_sublayer(
  77. conv_name,
  78. ConvBNLayer(
  79. ch_in=ch_in,
  80. ch_out=ch_out,
  81. filter_size=filter_size,
  82. padding=(filter_size - 1) // 2,
  83. norm_type=norm_type,
  84. freeze_norm=freeze_norm,
  85. data_format=data_format,
  86. name=name + post_name))
  87. self.tip = ConvBNLayer(
  88. ch_in=channel,
  89. ch_out=channel * 2,
  90. filter_size=3,
  91. padding=1,
  92. norm_type=norm_type,
  93. freeze_norm=freeze_norm,
  94. data_format=data_format,
  95. name=name + '.tip')
  96. def forward(self, inputs):
  97. route = self.conv_module(inputs)
  98. tip = self.tip(route)
  99. return route, tip
  100. class SPP(nn.Layer):
  101. def __init__(self,
  102. ch_in,
  103. ch_out,
  104. k,
  105. pool_size,
  106. norm_type='bn',
  107. freeze_norm=False,
  108. name='',
  109. act='leaky',
  110. data_format='NCHW'):
  111. """
  112. SPP layer, which consist of four pooling layer follwed by conv layer
  113. Args:
  114. ch_in (int): input channel of conv layer
  115. ch_out (int): output channel of conv layer
  116. k (int): kernel size of conv layer
  117. norm_type (str): batch norm type
  118. freeze_norm (bool): whether to freeze norm, default False
  119. name (str): layer name
  120. act (str): activation function
  121. data_format (str): data format, NCHW or NHWC
  122. """
  123. super(SPP, self).__init__()
  124. self.pool = []
  125. self.data_format = data_format
  126. for size in pool_size:
  127. pool = self.add_sublayer(
  128. '{}.pool1'.format(name),
  129. nn.MaxPool2D(
  130. kernel_size=size,
  131. stride=1,
  132. padding=size // 2,
  133. data_format=data_format,
  134. ceil_mode=False))
  135. self.pool.append(pool)
  136. self.conv = ConvBNLayer(
  137. ch_in,
  138. ch_out,
  139. k,
  140. padding=k // 2,
  141. norm_type=norm_type,
  142. freeze_norm=freeze_norm,
  143. name=name,
  144. act=act,
  145. data_format=data_format)
  146. def forward(self, x):
  147. outs = [x]
  148. for pool in self.pool:
  149. outs.append(pool(x))
  150. if self.data_format == "NCHW":
  151. y = paddle.concat(outs, axis=1)
  152. else:
  153. y = paddle.concat(outs, axis=-1)
  154. y = self.conv(y)
  155. return y
  156. class CoordConv(nn.Layer):
  157. def __init__(self,
  158. ch_in,
  159. ch_out,
  160. filter_size,
  161. padding,
  162. norm_type,
  163. freeze_norm=False,
  164. name='',
  165. data_format='NCHW'):
  166. """
  167. CoordConv layer, see https://arxiv.org/abs/1807.03247
  168. Args:
  169. ch_in (int): input channel
  170. ch_out (int): output channel
  171. filter_size (int): filter size, default 3
  172. padding (int): padding size, default 0
  173. norm_type (str): batch norm type, default bn
  174. name (str): layer name
  175. data_format (str): data format, NCHW or NHWC
  176. """
  177. super(CoordConv, self).__init__()
  178. self.conv = ConvBNLayer(
  179. ch_in + 2,
  180. ch_out,
  181. filter_size=filter_size,
  182. padding=padding,
  183. norm_type=norm_type,
  184. freeze_norm=freeze_norm,
  185. data_format=data_format,
  186. name=name)
  187. self.data_format = data_format
  188. def forward(self, x):
  189. gx, gy = add_coord(x, self.data_format)
  190. if self.data_format == 'NCHW':
  191. y = paddle.concat([x, gx, gy], axis=1)
  192. else:
  193. y = paddle.concat([x, gx, gy], axis=-1)
  194. y = self.conv(y)
  195. return y
  196. class PPYOLODetBlock(nn.Layer):
  197. def __init__(self, cfg, name, data_format='NCHW'):
  198. """
  199. PPYOLODetBlock layer
  200. Args:
  201. cfg (list): layer configs for this block
  202. name (str): block name
  203. data_format (str): data format, NCHW or NHWC
  204. """
  205. super(PPYOLODetBlock, self).__init__()
  206. self.conv_module = nn.Sequential()
  207. for idx, (conv_name, layer, args, kwargs) in enumerate(cfg[:-1]):
  208. kwargs.update(
  209. name='{}.{}'.format(name, conv_name), data_format=data_format)
  210. self.conv_module.add_sublayer(conv_name, layer(*args, **kwargs))
  211. conv_name, layer, args, kwargs = cfg[-1]
  212. kwargs.update(
  213. name='{}.{}'.format(name, conv_name), data_format=data_format)
  214. self.tip = layer(*args, **kwargs)
  215. def forward(self, inputs):
  216. route = self.conv_module(inputs)
  217. tip = self.tip(route)
  218. return route, tip
  219. class PPYOLOTinyDetBlock(nn.Layer):
  220. def __init__(self,
  221. ch_in,
  222. ch_out,
  223. name,
  224. drop_block=False,
  225. block_size=3,
  226. keep_prob=0.9,
  227. data_format='NCHW'):
  228. """
  229. PPYOLO Tiny DetBlock layer
  230. Args:
  231. ch_in (list): input channel number
  232. ch_out (list): output channel number
  233. name (str): block name
  234. drop_block: whether user DropBlock
  235. block_size: drop block size
  236. keep_prob: probability to keep block in DropBlock
  237. data_format (str): data format, NCHW or NHWC
  238. """
  239. super(PPYOLOTinyDetBlock, self).__init__()
  240. self.drop_block_ = drop_block
  241. self.conv_module = nn.Sequential()
  242. cfgs = [
  243. # name, in channels, out channels, filter_size,
  244. # stride, padding, groups
  245. ['.0', ch_in, ch_out, 1, 1, 0, 1],
  246. ['.1', ch_out, ch_out, 5, 1, 2, ch_out],
  247. ['.2', ch_out, ch_out, 1, 1, 0, 1],
  248. ['.route', ch_out, ch_out, 5, 1, 2, ch_out],
  249. ]
  250. for cfg in cfgs:
  251. conv_name, conv_ch_in, conv_ch_out, filter_size, stride, padding, \
  252. groups = cfg
  253. self.conv_module.add_sublayer(
  254. name + conv_name,
  255. ConvBNLayer(
  256. ch_in=conv_ch_in,
  257. ch_out=conv_ch_out,
  258. filter_size=filter_size,
  259. stride=stride,
  260. padding=padding,
  261. groups=groups,
  262. name=name + conv_name))
  263. self.tip = ConvBNLayer(
  264. ch_in=ch_out,
  265. ch_out=ch_out,
  266. filter_size=1,
  267. stride=1,
  268. padding=0,
  269. groups=1,
  270. name=name + conv_name)
  271. if self.drop_block_:
  272. self.drop_block = DropBlock(
  273. block_size=block_size,
  274. keep_prob=keep_prob,
  275. data_format=data_format,
  276. name=name + '.dropblock')
  277. def forward(self, inputs):
  278. if self.drop_block_:
  279. inputs = self.drop_block(inputs)
  280. route = self.conv_module(inputs)
  281. tip = self.tip(route)
  282. return route, tip
  283. class PPYOLODetBlockCSP(nn.Layer):
  284. def __init__(self,
  285. cfg,
  286. ch_in,
  287. ch_out,
  288. act,
  289. norm_type,
  290. name,
  291. data_format='NCHW'):
  292. """
  293. PPYOLODetBlockCSP layer
  294. Args:
  295. cfg (list): layer configs for this block
  296. ch_in (int): input channel
  297. ch_out (int): output channel
  298. act (str): default mish
  299. name (str): block name
  300. data_format (str): data format, NCHW or NHWC
  301. """
  302. super(PPYOLODetBlockCSP, self).__init__()
  303. self.data_format = data_format
  304. self.conv1 = ConvBNLayer(
  305. ch_in,
  306. ch_out,
  307. 1,
  308. padding=0,
  309. act=act,
  310. norm_type=norm_type,
  311. name=name + '.left',
  312. data_format=data_format)
  313. self.conv2 = ConvBNLayer(
  314. ch_in,
  315. ch_out,
  316. 1,
  317. padding=0,
  318. act=act,
  319. norm_type=norm_type,
  320. name=name + '.right',
  321. data_format=data_format)
  322. self.conv3 = ConvBNLayer(
  323. ch_out * 2,
  324. ch_out * 2,
  325. 1,
  326. padding=0,
  327. act=act,
  328. norm_type=norm_type,
  329. name=name,
  330. data_format=data_format)
  331. self.conv_module = nn.Sequential()
  332. for idx, (layer_name, layer, args, kwargs) in enumerate(cfg):
  333. kwargs.update(name=name + layer_name, data_format=data_format)
  334. self.conv_module.add_sublayer(layer_name, layer(*args, **kwargs))
  335. def forward(self, inputs):
  336. conv_left = self.conv1(inputs)
  337. conv_right = self.conv2(inputs)
  338. conv_left = self.conv_module(conv_left)
  339. if self.data_format == 'NCHW':
  340. conv = paddle.concat([conv_left, conv_right], axis=1)
  341. else:
  342. conv = paddle.concat([conv_left, conv_right], axis=-1)
  343. conv = self.conv3(conv)
  344. return conv, conv
  345. @register
  346. @serializable
  347. class YOLOv3FPN(nn.Layer):
  348. __shared__ = ['norm_type', 'data_format']
  349. def __init__(self,
  350. in_channels=[256, 512, 1024],
  351. norm_type='bn',
  352. freeze_norm=False,
  353. data_format='NCHW'):
  354. """
  355. YOLOv3FPN layer
  356. Args:
  357. in_channels (list): input channels for fpn
  358. norm_type (str): batch norm type, default bn
  359. data_format (str): data format, NCHW or NHWC
  360. """
  361. super(YOLOv3FPN, self).__init__()
  362. assert len(in_channels) > 0, "in_channels length should > 0"
  363. self.in_channels = in_channels
  364. self.num_blocks = len(in_channels)
  365. self._out_channels = []
  366. self.yolo_blocks = []
  367. self.routes = []
  368. self.data_format = data_format
  369. for i in range(self.num_blocks):
  370. name = 'yolo_block.{}'.format(i)
  371. in_channel = in_channels[-i - 1]
  372. if i > 0:
  373. in_channel += 512 // (2**i)
  374. yolo_block = self.add_sublayer(
  375. name,
  376. YoloDetBlock(
  377. in_channel,
  378. channel=512 // (2**i),
  379. norm_type=norm_type,
  380. freeze_norm=freeze_norm,
  381. data_format=data_format,
  382. name=name))
  383. self.yolo_blocks.append(yolo_block)
  384. # tip layer output channel doubled
  385. self._out_channels.append(1024 // (2**i))
  386. if i < self.num_blocks - 1:
  387. name = 'yolo_transition.{}'.format(i)
  388. route = self.add_sublayer(
  389. name,
  390. ConvBNLayer(
  391. ch_in=512 // (2**i),
  392. ch_out=256 // (2**i),
  393. filter_size=1,
  394. stride=1,
  395. padding=0,
  396. norm_type=norm_type,
  397. freeze_norm=freeze_norm,
  398. data_format=data_format,
  399. name=name))
  400. self.routes.append(route)
  401. def forward(self, blocks, for_mot=False):
  402. assert len(blocks) == self.num_blocks
  403. blocks = blocks[::-1]
  404. yolo_feats = []
  405. # add embedding features output for multi-object tracking model
  406. if for_mot:
  407. emb_feats = []
  408. for i, block in enumerate(blocks):
  409. if i > 0:
  410. if self.data_format == 'NCHW':
  411. block = paddle.concat([route, block], axis=1)
  412. else:
  413. block = paddle.concat([route, block], axis=-1)
  414. route, tip = self.yolo_blocks[i](block)
  415. yolo_feats.append(tip)
  416. if for_mot:
  417. # add embedding features output
  418. emb_feats.append(route)
  419. if i < self.num_blocks - 1:
  420. route = self.routes[i](route)
  421. route = F.interpolate(
  422. route, scale_factor=2., data_format=self.data_format)
  423. if for_mot:
  424. return {'yolo_feats': yolo_feats, 'emb_feats': emb_feats}
  425. else:
  426. return yolo_feats
  427. @classmethod
  428. def from_config(cls, cfg, input_shape):
  429. return {'in_channels': [i.channels for i in input_shape], }
  430. @property
  431. def out_shape(self):
  432. return [ShapeSpec(channels=c) for c in self._out_channels]
  433. @register
  434. @serializable
  435. class PPYOLOFPN(nn.Layer):
  436. __shared__ = ['norm_type', 'data_format']
  437. def __init__(self,
  438. in_channels=[512, 1024, 2048],
  439. norm_type='bn',
  440. freeze_norm=False,
  441. data_format='NCHW',
  442. coord_conv=False,
  443. conv_block_num=2,
  444. drop_block=False,
  445. block_size=3,
  446. keep_prob=0.9,
  447. spp=False):
  448. """
  449. PPYOLOFPN layer
  450. Args:
  451. in_channels (list): input channels for fpn
  452. norm_type (str): batch norm type, default bn
  453. data_format (str): data format, NCHW or NHWC
  454. coord_conv (bool): whether use CoordConv or not
  455. conv_block_num (int): conv block num of each pan block
  456. drop_block (bool): whether use DropBlock or not
  457. block_size (int): block size of DropBlock
  458. keep_prob (float): keep probability of DropBlock
  459. spp (bool): whether use spp or not
  460. """
  461. super(PPYOLOFPN, self).__init__()
  462. assert len(in_channels) > 0, "in_channels length should > 0"
  463. self.in_channels = in_channels
  464. self.num_blocks = len(in_channels)
  465. # parse kwargs
  466. self.coord_conv = coord_conv
  467. self.drop_block = drop_block
  468. self.block_size = block_size
  469. self.keep_prob = keep_prob
  470. self.spp = spp
  471. self.conv_block_num = conv_block_num
  472. self.data_format = data_format
  473. if self.coord_conv:
  474. ConvLayer = CoordConv
  475. else:
  476. ConvLayer = ConvBNLayer
  477. if self.drop_block:
  478. dropblock_cfg = [[
  479. 'dropblock', DropBlock, [self.block_size, self.keep_prob],
  480. dict()
  481. ]]
  482. else:
  483. dropblock_cfg = []
  484. self._out_channels = []
  485. self.yolo_blocks = []
  486. self.routes = []
  487. for i, ch_in in enumerate(self.in_channels[::-1]):
  488. if i > 0:
  489. ch_in += 512 // (2**i)
  490. channel = 64 * (2**self.num_blocks) // (2**i)
  491. base_cfg = []
  492. c_in, c_out = ch_in, channel
  493. for j in range(self.conv_block_num):
  494. base_cfg += [
  495. [
  496. 'conv{}'.format(2 * j), ConvLayer, [c_in, c_out, 1],
  497. dict(
  498. padding=0,
  499. norm_type=norm_type,
  500. freeze_norm=freeze_norm)
  501. ],
  502. [
  503. 'conv{}'.format(2 * j + 1), ConvBNLayer,
  504. [c_out, c_out * 2, 3], dict(
  505. padding=1,
  506. norm_type=norm_type,
  507. freeze_norm=freeze_norm)
  508. ],
  509. ]
  510. c_in, c_out = c_out * 2, c_out
  511. base_cfg += [[
  512. 'route', ConvLayer, [c_in, c_out, 1], dict(
  513. padding=0, norm_type=norm_type, freeze_norm=freeze_norm)
  514. ], [
  515. 'tip', ConvLayer, [c_out, c_out * 2, 3], dict(
  516. padding=1, norm_type=norm_type, freeze_norm=freeze_norm)
  517. ]]
  518. if self.conv_block_num == 2:
  519. if i == 0:
  520. if self.spp:
  521. spp_cfg = [[
  522. 'spp', SPP, [channel * 4, channel, 1], dict(
  523. pool_size=[5, 9, 13],
  524. norm_type=norm_type,
  525. freeze_norm=freeze_norm)
  526. ]]
  527. else:
  528. spp_cfg = []
  529. cfg = base_cfg[0:3] + spp_cfg + base_cfg[
  530. 3:4] + dropblock_cfg + base_cfg[4:6]
  531. else:
  532. cfg = base_cfg[0:2] + dropblock_cfg + base_cfg[2:6]
  533. elif self.conv_block_num == 0:
  534. if self.spp and i == 0:
  535. spp_cfg = [[
  536. 'spp', SPP, [c_in * 4, c_in, 1], dict(
  537. pool_size=[5, 9, 13],
  538. norm_type=norm_type,
  539. freeze_norm=freeze_norm)
  540. ]]
  541. else:
  542. spp_cfg = []
  543. cfg = spp_cfg + dropblock_cfg + base_cfg
  544. name = 'yolo_block.{}'.format(i)
  545. yolo_block = self.add_sublayer(name, PPYOLODetBlock(cfg, name))
  546. self.yolo_blocks.append(yolo_block)
  547. self._out_channels.append(channel * 2)
  548. if i < self.num_blocks - 1:
  549. name = 'yolo_transition.{}'.format(i)
  550. route = self.add_sublayer(
  551. name,
  552. ConvBNLayer(
  553. ch_in=channel,
  554. ch_out=256 // (2**i),
  555. filter_size=1,
  556. stride=1,
  557. padding=0,
  558. norm_type=norm_type,
  559. freeze_norm=freeze_norm,
  560. data_format=data_format,
  561. name=name))
  562. self.routes.append(route)
  563. def forward(self, blocks, for_mot=False):
  564. assert len(blocks) == self.num_blocks
  565. blocks = blocks[::-1]
  566. yolo_feats = []
  567. # add embedding features output for multi-object tracking model
  568. if for_mot:
  569. emb_feats = []
  570. for i, block in enumerate(blocks):
  571. if i > 0:
  572. if self.data_format == 'NCHW':
  573. block = paddle.concat([route, block], axis=1)
  574. else:
  575. block = paddle.concat([route, block], axis=-1)
  576. route, tip = self.yolo_blocks[i](block)
  577. yolo_feats.append(tip)
  578. if for_mot:
  579. # add embedding features output
  580. emb_feats.append(route)
  581. if i < self.num_blocks - 1:
  582. route = self.routes[i](route)
  583. route = F.interpolate(
  584. route, scale_factor=2., data_format=self.data_format)
  585. if for_mot:
  586. return {'yolo_feats': yolo_feats, 'emb_feats': emb_feats}
  587. else:
  588. return yolo_feats
  589. @classmethod
  590. def from_config(cls, cfg, input_shape):
  591. return {'in_channels': [i.channels for i in input_shape], }
  592. @property
  593. def out_shape(self):
  594. return [ShapeSpec(channels=c) for c in self._out_channels]
  595. @register
  596. @serializable
  597. class PPYOLOTinyFPN(nn.Layer):
  598. __shared__ = ['norm_type', 'data_format']
  599. def __init__(self,
  600. in_channels=[80, 56, 34],
  601. detection_block_channels=[160, 128, 96],
  602. norm_type='bn',
  603. data_format='NCHW',
  604. **kwargs):
  605. """
  606. PPYOLO Tiny FPN layer
  607. Args:
  608. in_channels (list): input channels for fpn
  609. detection_block_channels (list): channels in fpn
  610. norm_type (str): batch norm type, default bn
  611. data_format (str): data format, NCHW or NHWC
  612. kwargs: extra key-value pairs, such as parameter of DropBlock and spp
  613. """
  614. super(PPYOLOTinyFPN, self).__init__()
  615. assert len(in_channels) > 0, "in_channels length should > 0"
  616. self.in_channels = in_channels[::-1]
  617. assert len(detection_block_channels
  618. ) > 0, "detection_block_channelslength should > 0"
  619. self.detection_block_channels = detection_block_channels
  620. self.data_format = data_format
  621. self.num_blocks = len(in_channels)
  622. # parse kwargs
  623. self.drop_block = kwargs.get('drop_block', False)
  624. self.block_size = kwargs.get('block_size', 3)
  625. self.keep_prob = kwargs.get('keep_prob', 0.9)
  626. self.spp_ = kwargs.get('spp', False)
  627. if self.spp_:
  628. self.spp = SPP(self.in_channels[0] * 4,
  629. self.in_channels[0],
  630. k=1,
  631. pool_size=[5, 9, 13],
  632. norm_type=norm_type,
  633. name='spp')
  634. self._out_channels = []
  635. self.yolo_blocks = []
  636. self.routes = []
  637. for i, (
  638. ch_in, ch_out
  639. ) in enumerate(zip(self.in_channels, self.detection_block_channels)):
  640. name = 'yolo_block.{}'.format(i)
  641. if i > 0:
  642. ch_in += self.detection_block_channels[i - 1]
  643. yolo_block = self.add_sublayer(
  644. name,
  645. PPYOLOTinyDetBlock(
  646. ch_in,
  647. ch_out,
  648. name,
  649. drop_block=self.drop_block,
  650. block_size=self.block_size,
  651. keep_prob=self.keep_prob))
  652. self.yolo_blocks.append(yolo_block)
  653. self._out_channels.append(ch_out)
  654. if i < self.num_blocks - 1:
  655. name = 'yolo_transition.{}'.format(i)
  656. route = self.add_sublayer(
  657. name,
  658. ConvBNLayer(
  659. ch_in=ch_out,
  660. ch_out=ch_out,
  661. filter_size=1,
  662. stride=1,
  663. padding=0,
  664. norm_type=norm_type,
  665. data_format=data_format,
  666. name=name))
  667. self.routes.append(route)
  668. def forward(self, blocks, for_mot=False):
  669. assert len(blocks) == self.num_blocks
  670. blocks = blocks[::-1]
  671. yolo_feats = []
  672. # add embedding features output for multi-object tracking model
  673. if for_mot:
  674. emb_feats = []
  675. for i, block in enumerate(blocks):
  676. if i == 0 and self.spp_:
  677. block = self.spp(block)
  678. if i > 0:
  679. if self.data_format == 'NCHW':
  680. block = paddle.concat([route, block], axis=1)
  681. else:
  682. block = paddle.concat([route, block], axis=-1)
  683. route, tip = self.yolo_blocks[i](block)
  684. yolo_feats.append(tip)
  685. if for_mot:
  686. # add embedding features output
  687. emb_feats.append(route)
  688. if i < self.num_blocks - 1:
  689. route = self.routes[i](route)
  690. route = F.interpolate(
  691. route, scale_factor=2., data_format=self.data_format)
  692. if for_mot:
  693. return {'yolo_feats': yolo_feats, 'emb_feats': emb_feats}
  694. else:
  695. return yolo_feats
  696. @classmethod
  697. def from_config(cls, cfg, input_shape):
  698. return {'in_channels': [i.channels for i in input_shape], }
  699. @property
  700. def out_shape(self):
  701. return [ShapeSpec(channels=c) for c in self._out_channels]
  702. @register
  703. @serializable
  704. class PPYOLOPAN(nn.Layer):
  705. __shared__ = ['norm_type', 'data_format']
  706. def __init__(self,
  707. in_channels=[512, 1024, 2048],
  708. norm_type='bn',
  709. data_format='NCHW',
  710. act='mish',
  711. conv_block_num=3,
  712. drop_block=False,
  713. block_size=3,
  714. keep_prob=0.9,
  715. spp=False):
  716. """
  717. PPYOLOPAN layer with SPP, DropBlock and CSP connection.
  718. Args:
  719. in_channels (list): input channels for fpn
  720. norm_type (str): batch norm type, default bn
  721. data_format (str): data format, NCHW or NHWC
  722. act (str): activation function, default mish
  723. conv_block_num (int): conv block num of each pan block
  724. drop_block (bool): whether use DropBlock or not
  725. block_size (int): block size of DropBlock
  726. keep_prob (float): keep probability of DropBlock
  727. spp (bool): whether use spp or not
  728. """
  729. super(PPYOLOPAN, self).__init__()
  730. assert len(in_channels) > 0, "in_channels length should > 0"
  731. self.in_channels = in_channels
  732. self.num_blocks = len(in_channels)
  733. # parse kwargs
  734. self.drop_block = drop_block
  735. self.block_size = block_size
  736. self.keep_prob = keep_prob
  737. self.spp = spp
  738. self.conv_block_num = conv_block_num
  739. self.data_format = data_format
  740. if self.drop_block:
  741. dropblock_cfg = [[
  742. 'dropblock', DropBlock, [self.block_size, self.keep_prob],
  743. dict()
  744. ]]
  745. else:
  746. dropblock_cfg = []
  747. # fpn
  748. self.fpn_blocks = []
  749. self.fpn_routes = []
  750. fpn_channels = []
  751. for i, ch_in in enumerate(self.in_channels[::-1]):
  752. if i > 0:
  753. ch_in += 512 // (2**(i - 1))
  754. channel = 512 // (2**i)
  755. base_cfg = []
  756. for j in range(self.conv_block_num):
  757. base_cfg += [
  758. # name, layer, args
  759. [
  760. '{}.0'.format(j), ConvBNLayer, [channel, channel, 1],
  761. dict(
  762. padding=0, act=act, norm_type=norm_type)
  763. ],
  764. [
  765. '{}.1'.format(j), ConvBNLayer, [channel, channel, 3],
  766. dict(
  767. padding=1, act=act, norm_type=norm_type)
  768. ]
  769. ]
  770. if i == 0 and self.spp:
  771. base_cfg[3] = [
  772. 'spp', SPP, [channel * 4, channel, 1], dict(
  773. pool_size=[5, 9, 13], act=act, norm_type=norm_type)
  774. ]
  775. cfg = base_cfg[:4] + dropblock_cfg + base_cfg[4:]
  776. name = 'fpn.{}'.format(i)
  777. fpn_block = self.add_sublayer(
  778. name,
  779. PPYOLODetBlockCSP(cfg, ch_in, channel, act, norm_type, name,
  780. data_format))
  781. self.fpn_blocks.append(fpn_block)
  782. fpn_channels.append(channel * 2)
  783. if i < self.num_blocks - 1:
  784. name = 'fpn_transition.{}'.format(i)
  785. route = self.add_sublayer(
  786. name,
  787. ConvBNLayer(
  788. ch_in=channel * 2,
  789. ch_out=channel,
  790. filter_size=1,
  791. stride=1,
  792. padding=0,
  793. act=act,
  794. norm_type=norm_type,
  795. data_format=data_format,
  796. name=name))
  797. self.fpn_routes.append(route)
  798. # pan
  799. self.pan_blocks = []
  800. self.pan_routes = []
  801. self._out_channels = [512 // (2**(self.num_blocks - 2)), ]
  802. for i in reversed(range(self.num_blocks - 1)):
  803. name = 'pan_transition.{}'.format(i)
  804. route = self.add_sublayer(
  805. name,
  806. ConvBNLayer(
  807. ch_in=fpn_channels[i + 1],
  808. ch_out=fpn_channels[i + 1],
  809. filter_size=3,
  810. stride=2,
  811. padding=1,
  812. act=act,
  813. norm_type=norm_type,
  814. data_format=data_format,
  815. name=name))
  816. self.pan_routes = [route, ] + self.pan_routes
  817. base_cfg = []
  818. ch_in = fpn_channels[i] + fpn_channels[i + 1]
  819. channel = 512 // (2**i)
  820. for j in range(self.conv_block_num):
  821. base_cfg += [
  822. # name, layer, args
  823. [
  824. '{}.0'.format(j), ConvBNLayer, [channel, channel, 1],
  825. dict(
  826. padding=0, act=act, norm_type=norm_type)
  827. ],
  828. [
  829. '{}.1'.format(j), ConvBNLayer, [channel, channel, 3],
  830. dict(
  831. padding=1, act=act, norm_type=norm_type)
  832. ]
  833. ]
  834. cfg = base_cfg[:4] + dropblock_cfg + base_cfg[4:]
  835. name = 'pan.{}'.format(i)
  836. pan_block = self.add_sublayer(
  837. name,
  838. PPYOLODetBlockCSP(cfg, ch_in, channel, act, norm_type, name,
  839. data_format))
  840. self.pan_blocks = [pan_block, ] + self.pan_blocks
  841. self._out_channels.append(channel * 2)
  842. self._out_channels = self._out_channels[::-1]
  843. def forward(self, blocks, for_mot=False):
  844. assert len(blocks) == self.num_blocks
  845. blocks = blocks[::-1]
  846. fpn_feats = []
  847. # add embedding features output for multi-object tracking model
  848. if for_mot:
  849. emb_feats = []
  850. for i, block in enumerate(blocks):
  851. if i > 0:
  852. if self.data_format == 'NCHW':
  853. block = paddle.concat([route, block], axis=1)
  854. else:
  855. block = paddle.concat([route, block], axis=-1)
  856. route, tip = self.fpn_blocks[i](block)
  857. fpn_feats.append(tip)
  858. if for_mot:
  859. # add embedding features output
  860. emb_feats.append(route)
  861. if i < self.num_blocks - 1:
  862. route = self.fpn_routes[i](route)
  863. route = F.interpolate(
  864. route, scale_factor=2., data_format=self.data_format)
  865. pan_feats = [fpn_feats[-1], ]
  866. route = fpn_feats[self.num_blocks - 1]
  867. for i in reversed(range(self.num_blocks - 1)):
  868. block = fpn_feats[i]
  869. route = self.pan_routes[i](route)
  870. if self.data_format == 'NCHW':
  871. block = paddle.concat([route, block], axis=1)
  872. else:
  873. block = paddle.concat([route, block], axis=-1)
  874. route, tip = self.pan_blocks[i](block)
  875. pan_feats.append(tip)
  876. if for_mot:
  877. return {'yolo_feats': pan_feats[::-1], 'emb_feats': emb_feats}
  878. else:
  879. return pan_feats[::-1]
  880. @classmethod
  881. def from_config(cls, cfg, input_shape):
  882. return {'in_channels': [i.channels for i in input_shape], }
  883. @property
  884. def out_shape(self):
  885. return [ShapeSpec(channels=c) for c in self._out_channels]
  886. @register
  887. @serializable
  888. class YOLOCSPPAN(nn.Layer):
  889. """
  890. YOLO CSP-PAN, used in YOLOv5 and YOLOX.
  891. """
  892. __shared__ = ['depth_mult', 'data_format', 'act', 'trt']
  893. def __init__(self,
  894. depth_mult=1.0,
  895. in_channels=[256, 512, 1024],
  896. depthwise=False,
  897. data_format='NCHW',
  898. act='silu',
  899. trt=False):
  900. super(YOLOCSPPAN, self).__init__()
  901. self.in_channels = in_channels
  902. self._out_channels = in_channels
  903. Conv = DWConv if depthwise else BaseConv
  904. self.data_format = data_format
  905. act = get_act_fn(
  906. act, trt=trt) if act is None or isinstance(act,
  907. (str, dict)) else act
  908. self.upsample = nn.Upsample(scale_factor=2, mode="nearest")
  909. # top-down fpn
  910. self.lateral_convs = nn.LayerList()
  911. self.fpn_blocks = nn.LayerList()
  912. for idx in range(len(in_channels) - 1, 0, -1):
  913. self.lateral_convs.append(
  914. BaseConv(
  915. int(in_channels[idx]),
  916. int(in_channels[idx - 1]),
  917. 1,
  918. 1,
  919. act=act))
  920. self.fpn_blocks.append(
  921. CSPLayer(
  922. int(in_channels[idx - 1] * 2),
  923. int(in_channels[idx - 1]),
  924. round(3 * depth_mult),
  925. shortcut=False,
  926. depthwise=depthwise,
  927. act=act))
  928. # bottom-up pan
  929. self.downsample_convs = nn.LayerList()
  930. self.pan_blocks = nn.LayerList()
  931. for idx in range(len(in_channels) - 1):
  932. self.downsample_convs.append(
  933. Conv(
  934. int(in_channels[idx]),
  935. int(in_channels[idx]),
  936. 3,
  937. stride=2,
  938. act=act))
  939. self.pan_blocks.append(
  940. CSPLayer(
  941. int(in_channels[idx] * 2),
  942. int(in_channels[idx + 1]),
  943. round(3 * depth_mult),
  944. shortcut=False,
  945. depthwise=depthwise,
  946. act=act))
  947. def forward(self, feats, for_mot=False):
  948. assert len(feats) == len(self.in_channels)
  949. # top-down fpn
  950. inner_outs = [feats[-1]]
  951. for idx in range(len(self.in_channels) - 1, 0, -1):
  952. feat_heigh = inner_outs[0]
  953. feat_low = feats[idx - 1]
  954. feat_heigh = self.lateral_convs[len(self.in_channels) - 1 - idx](
  955. feat_heigh)
  956. inner_outs[0] = feat_heigh
  957. upsample_feat = F.interpolate(
  958. feat_heigh,
  959. scale_factor=2.,
  960. mode="nearest",
  961. data_format=self.data_format)
  962. inner_out = self.fpn_blocks[len(self.in_channels) - 1 - idx](
  963. paddle.concat(
  964. [upsample_feat, feat_low], axis=1))
  965. inner_outs.insert(0, inner_out)
  966. # bottom-up pan
  967. outs = [inner_outs[0]]
  968. for idx in range(len(self.in_channels) - 1):
  969. feat_low = outs[-1]
  970. feat_height = inner_outs[idx + 1]
  971. downsample_feat = self.downsample_convs[idx](feat_low)
  972. out = self.pan_blocks[idx](paddle.concat(
  973. [downsample_feat, feat_height], axis=1))
  974. outs.append(out)
  975. return outs
  976. @classmethod
  977. def from_config(cls, cfg, input_shape):
  978. return {'in_channels': [i.channels for i in input_shape], }
  979. @property
  980. def out_shape(self):
  981. return [ShapeSpec(channels=c) for c in self._out_channels]