hrnet.py 24 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555556557558559560561562563564565566567568569570571572573574575576577578579580581582583584585586587588589590591592593594595596597598599600601602603604605606607608609610611612613614615616617618619620621622623624625626627628629630631632633634635636637638639640641642643644645646647648649650651652653654655656657658659660661662663664665666667668669670671672673674675676677678679680681682683684685686687688689690691692693694695696697698699700701702703704705706707708709710711712713714715716717718719720721722723724
  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 paddle.regularizer import L2Decay
  18. from paddle import ParamAttr
  19. from paddle.nn.initializer import Normal
  20. from numbers import Integral
  21. import math
  22. from paddlex.ppdet.core.workspace import register, serializable
  23. from ..shape_spec import ShapeSpec
  24. __all__ = ['HRNet']
  25. class ConvNormLayer(nn.Layer):
  26. def __init__(self,
  27. ch_in,
  28. ch_out,
  29. filter_size,
  30. stride=1,
  31. norm_type='bn',
  32. norm_groups=32,
  33. use_dcn=False,
  34. norm_decay=0.,
  35. freeze_norm=False,
  36. act=None,
  37. name=None):
  38. super(ConvNormLayer, self).__init__()
  39. assert norm_type in ['bn', 'sync_bn', 'gn']
  40. self.act = act
  41. self.conv = nn.Conv2D(
  42. in_channels=ch_in,
  43. out_channels=ch_out,
  44. kernel_size=filter_size,
  45. stride=stride,
  46. padding=(filter_size - 1) // 2,
  47. groups=1,
  48. weight_attr=ParamAttr(
  49. name=name + "_weights", initializer=Normal(
  50. mean=0., std=0.01)),
  51. bias_attr=False)
  52. norm_lr = 0. if freeze_norm else 1.
  53. norm_name = name + '_bn'
  54. param_attr = ParamAttr(
  55. name=norm_name + "_scale",
  56. learning_rate=norm_lr,
  57. regularizer=L2Decay(norm_decay))
  58. bias_attr = ParamAttr(
  59. name=norm_name + "_offset",
  60. learning_rate=norm_lr,
  61. regularizer=L2Decay(norm_decay))
  62. global_stats = True if freeze_norm else False
  63. if norm_type in ['bn', 'sync_bn']:
  64. self.norm = nn.BatchNorm(
  65. ch_out,
  66. param_attr=param_attr,
  67. bias_attr=bias_attr,
  68. use_global_stats=global_stats,
  69. moving_mean_name=norm_name + '_mean',
  70. moving_variance_name=norm_name + '_variance')
  71. elif norm_type == 'gn':
  72. self.norm = nn.GroupNorm(
  73. num_groups=norm_groups,
  74. num_channels=ch_out,
  75. weight_attr=param_attr,
  76. bias_attr=bias_attr)
  77. norm_params = self.norm.parameters()
  78. if freeze_norm:
  79. for param in norm_params:
  80. param.stop_gradient = True
  81. def forward(self, inputs):
  82. out = self.conv(inputs)
  83. out = self.norm(out)
  84. if self.act == 'relu':
  85. out = F.relu(out)
  86. return out
  87. class Layer1(nn.Layer):
  88. def __init__(self,
  89. num_channels,
  90. has_se=False,
  91. norm_decay=0.,
  92. freeze_norm=True,
  93. name=None):
  94. super(Layer1, self).__init__()
  95. self.bottleneck_block_list = []
  96. for i in range(4):
  97. bottleneck_block = self.add_sublayer(
  98. "block_{}_{}".format(name, i + 1),
  99. BottleneckBlock(
  100. num_channels=num_channels if i == 0 else 256,
  101. num_filters=64,
  102. has_se=has_se,
  103. stride=1,
  104. downsample=True if i == 0 else False,
  105. norm_decay=norm_decay,
  106. freeze_norm=freeze_norm,
  107. name=name + '_' + str(i + 1)))
  108. self.bottleneck_block_list.append(bottleneck_block)
  109. def forward(self, input):
  110. conv = input
  111. for block_func in self.bottleneck_block_list:
  112. conv = block_func(conv)
  113. return conv
  114. class TransitionLayer(nn.Layer):
  115. def __init__(self,
  116. in_channels,
  117. out_channels,
  118. norm_decay=0.,
  119. freeze_norm=True,
  120. name=None):
  121. super(TransitionLayer, self).__init__()
  122. num_in = len(in_channels)
  123. num_out = len(out_channels)
  124. out = []
  125. self.conv_bn_func_list = []
  126. for i in range(num_out):
  127. residual = None
  128. if i < num_in:
  129. if in_channels[i] != out_channels[i]:
  130. residual = self.add_sublayer(
  131. "transition_{}_layer_{}".format(name, i + 1),
  132. ConvNormLayer(
  133. ch_in=in_channels[i],
  134. ch_out=out_channels[i],
  135. filter_size=3,
  136. norm_decay=norm_decay,
  137. freeze_norm=freeze_norm,
  138. act='relu',
  139. name=name + '_layer_' + str(i + 1)))
  140. else:
  141. residual = self.add_sublayer(
  142. "transition_{}_layer_{}".format(name, i + 1),
  143. ConvNormLayer(
  144. ch_in=in_channels[-1],
  145. ch_out=out_channels[i],
  146. filter_size=3,
  147. stride=2,
  148. norm_decay=norm_decay,
  149. freeze_norm=freeze_norm,
  150. act='relu',
  151. name=name + '_layer_' + str(i + 1)))
  152. self.conv_bn_func_list.append(residual)
  153. def forward(self, input):
  154. outs = []
  155. for idx, conv_bn_func in enumerate(self.conv_bn_func_list):
  156. if conv_bn_func is None:
  157. outs.append(input[idx])
  158. else:
  159. if idx < len(input):
  160. outs.append(conv_bn_func(input[idx]))
  161. else:
  162. outs.append(conv_bn_func(input[-1]))
  163. return outs
  164. class Branches(nn.Layer):
  165. def __init__(self,
  166. block_num,
  167. in_channels,
  168. out_channels,
  169. has_se=False,
  170. norm_decay=0.,
  171. freeze_norm=True,
  172. name=None):
  173. super(Branches, self).__init__()
  174. self.basic_block_list = []
  175. for i in range(len(out_channels)):
  176. self.basic_block_list.append([])
  177. for j in range(block_num):
  178. in_ch = in_channels[i] if j == 0 else out_channels[i]
  179. basic_block_func = self.add_sublayer(
  180. "bb_{}_branch_layer_{}_{}".format(name, i + 1, j + 1),
  181. BasicBlock(
  182. num_channels=in_ch,
  183. num_filters=out_channels[i],
  184. has_se=has_se,
  185. norm_decay=norm_decay,
  186. freeze_norm=freeze_norm,
  187. name=name + '_branch_layer_' + str(i + 1) + '_' +
  188. str(j + 1)))
  189. self.basic_block_list[i].append(basic_block_func)
  190. def forward(self, inputs):
  191. outs = []
  192. for idx, input in enumerate(inputs):
  193. conv = input
  194. basic_block_list = self.basic_block_list[idx]
  195. for basic_block_func in basic_block_list:
  196. conv = basic_block_func(conv)
  197. outs.append(conv)
  198. return outs
  199. class BottleneckBlock(nn.Layer):
  200. def __init__(self,
  201. num_channels,
  202. num_filters,
  203. has_se,
  204. stride=1,
  205. downsample=False,
  206. norm_decay=0.,
  207. freeze_norm=True,
  208. name=None):
  209. super(BottleneckBlock, self).__init__()
  210. self.has_se = has_se
  211. self.downsample = downsample
  212. self.conv1 = ConvNormLayer(
  213. ch_in=num_channels,
  214. ch_out=num_filters,
  215. filter_size=1,
  216. norm_decay=norm_decay,
  217. freeze_norm=freeze_norm,
  218. act="relu",
  219. name=name + "_conv1")
  220. self.conv2 = ConvNormLayer(
  221. ch_in=num_filters,
  222. ch_out=num_filters,
  223. filter_size=3,
  224. stride=stride,
  225. norm_decay=norm_decay,
  226. freeze_norm=freeze_norm,
  227. act="relu",
  228. name=name + "_conv2")
  229. self.conv3 = ConvNormLayer(
  230. ch_in=num_filters,
  231. ch_out=num_filters * 4,
  232. filter_size=1,
  233. norm_decay=norm_decay,
  234. freeze_norm=freeze_norm,
  235. act=None,
  236. name=name + "_conv3")
  237. if self.downsample:
  238. self.conv_down = ConvNormLayer(
  239. ch_in=num_channels,
  240. ch_out=num_filters * 4,
  241. filter_size=1,
  242. norm_decay=norm_decay,
  243. freeze_norm=freeze_norm,
  244. act=None,
  245. name=name + "_downsample")
  246. if self.has_se:
  247. self.se = SELayer(
  248. num_channels=num_filters * 4,
  249. num_filters=num_filters * 4,
  250. reduction_ratio=16,
  251. name='fc' + name)
  252. def forward(self, input):
  253. residual = input
  254. conv1 = self.conv1(input)
  255. conv2 = self.conv2(conv1)
  256. conv3 = self.conv3(conv2)
  257. if self.downsample:
  258. residual = self.conv_down(input)
  259. if self.has_se:
  260. conv3 = self.se(conv3)
  261. y = paddle.add(x=residual, y=conv3)
  262. y = F.relu(y)
  263. return y
  264. class BasicBlock(nn.Layer):
  265. def __init__(self,
  266. num_channels,
  267. num_filters,
  268. stride=1,
  269. has_se=False,
  270. downsample=False,
  271. norm_decay=0.,
  272. freeze_norm=True,
  273. name=None):
  274. super(BasicBlock, self).__init__()
  275. self.has_se = has_se
  276. self.downsample = downsample
  277. self.conv1 = ConvNormLayer(
  278. ch_in=num_channels,
  279. ch_out=num_filters,
  280. filter_size=3,
  281. norm_decay=norm_decay,
  282. freeze_norm=freeze_norm,
  283. stride=stride,
  284. act="relu",
  285. name=name + "_conv1")
  286. self.conv2 = ConvNormLayer(
  287. ch_in=num_filters,
  288. ch_out=num_filters,
  289. filter_size=3,
  290. norm_decay=norm_decay,
  291. freeze_norm=freeze_norm,
  292. stride=1,
  293. act=None,
  294. name=name + "_conv2")
  295. if self.downsample:
  296. self.conv_down = ConvNormLayer(
  297. ch_in=num_channels,
  298. ch_out=num_filters * 4,
  299. filter_size=1,
  300. norm_decay=norm_decay,
  301. freeze_norm=freeze_norm,
  302. act=None,
  303. name=name + "_downsample")
  304. if self.has_se:
  305. self.se = SELayer(
  306. num_channels=num_filters,
  307. num_filters=num_filters,
  308. reduction_ratio=16,
  309. name='fc' + name)
  310. def forward(self, input):
  311. residual = input
  312. conv1 = self.conv1(input)
  313. conv2 = self.conv2(conv1)
  314. if self.downsample:
  315. residual = self.conv_down(input)
  316. if self.has_se:
  317. conv2 = self.se(conv2)
  318. y = paddle.add(x=residual, y=conv2)
  319. y = F.relu(y)
  320. return y
  321. class SELayer(nn.Layer):
  322. def __init__(self, num_channels, num_filters, reduction_ratio, name=None):
  323. super(SELayer, self).__init__()
  324. self.pool2d_gap = AdaptiveAvgPool2D(1)
  325. self._num_channels = num_channels
  326. med_ch = int(num_channels / reduction_ratio)
  327. stdv = 1.0 / math.sqrt(num_channels * 1.0)
  328. self.squeeze = Linear(
  329. num_channels,
  330. med_ch,
  331. weight_attr=ParamAttr(
  332. initializer=Uniform(-stdv, stdv), name=name + "_sqz_weights"),
  333. bias_attr=ParamAttr(name=name + '_sqz_offset'))
  334. stdv = 1.0 / math.sqrt(med_ch * 1.0)
  335. self.excitation = Linear(
  336. med_ch,
  337. num_filters,
  338. weight_attr=ParamAttr(
  339. initializer=Uniform(-stdv, stdv), name=name + "_exc_weights"),
  340. bias_attr=ParamAttr(name=name + '_exc_offset'))
  341. def forward(self, input):
  342. pool = self.pool2d_gap(input)
  343. pool = paddle.squeeze(pool, axis=[2, 3])
  344. squeeze = self.squeeze(pool)
  345. squeeze = F.relu(squeeze)
  346. excitation = self.excitation(squeeze)
  347. excitation = F.sigmoid(excitation)
  348. excitation = paddle.unsqueeze(excitation, axis=[2, 3])
  349. out = input * excitation
  350. return out
  351. class Stage(nn.Layer):
  352. def __init__(self,
  353. num_channels,
  354. num_modules,
  355. num_filters,
  356. has_se=False,
  357. norm_decay=0.,
  358. freeze_norm=True,
  359. multi_scale_output=True,
  360. name=None):
  361. super(Stage, self).__init__()
  362. self._num_modules = num_modules
  363. self.stage_func_list = []
  364. for i in range(num_modules):
  365. if i == num_modules - 1 and not multi_scale_output:
  366. stage_func = self.add_sublayer(
  367. "stage_{}_{}".format(name, i + 1),
  368. HighResolutionModule(
  369. num_channels=num_channels,
  370. num_filters=num_filters,
  371. has_se=has_se,
  372. norm_decay=norm_decay,
  373. freeze_norm=freeze_norm,
  374. multi_scale_output=False,
  375. name=name + '_' + str(i + 1)))
  376. else:
  377. stage_func = self.add_sublayer(
  378. "stage_{}_{}".format(name, i + 1),
  379. HighResolutionModule(
  380. num_channels=num_channels,
  381. num_filters=num_filters,
  382. has_se=has_se,
  383. norm_decay=norm_decay,
  384. freeze_norm=freeze_norm,
  385. name=name + '_' + str(i + 1)))
  386. self.stage_func_list.append(stage_func)
  387. def forward(self, input):
  388. out = input
  389. for idx in range(self._num_modules):
  390. out = self.stage_func_list[idx](out)
  391. return out
  392. class HighResolutionModule(nn.Layer):
  393. def __init__(self,
  394. num_channels,
  395. num_filters,
  396. has_se=False,
  397. multi_scale_output=True,
  398. norm_decay=0.,
  399. freeze_norm=True,
  400. name=None):
  401. super(HighResolutionModule, self).__init__()
  402. self.branches_func = Branches(
  403. block_num=4,
  404. in_channels=num_channels,
  405. out_channels=num_filters,
  406. has_se=has_se,
  407. norm_decay=norm_decay,
  408. freeze_norm=freeze_norm,
  409. name=name)
  410. self.fuse_func = FuseLayers(
  411. in_channels=num_filters,
  412. out_channels=num_filters,
  413. multi_scale_output=multi_scale_output,
  414. norm_decay=norm_decay,
  415. freeze_norm=freeze_norm,
  416. name=name)
  417. def forward(self, input):
  418. out = self.branches_func(input)
  419. out = self.fuse_func(out)
  420. return out
  421. class FuseLayers(nn.Layer):
  422. def __init__(self,
  423. in_channels,
  424. out_channels,
  425. multi_scale_output=True,
  426. norm_decay=0.,
  427. freeze_norm=True,
  428. name=None):
  429. super(FuseLayers, self).__init__()
  430. self._actual_ch = len(in_channels) if multi_scale_output else 1
  431. self._in_channels = in_channels
  432. self.residual_func_list = []
  433. for i in range(self._actual_ch):
  434. for j in range(len(in_channels)):
  435. residual_func = None
  436. if j > i:
  437. residual_func = self.add_sublayer(
  438. "residual_{}_layer_{}_{}".format(name, i + 1, j + 1),
  439. ConvNormLayer(
  440. ch_in=in_channels[j],
  441. ch_out=out_channels[i],
  442. filter_size=1,
  443. stride=1,
  444. act=None,
  445. norm_decay=norm_decay,
  446. freeze_norm=freeze_norm,
  447. name=name + '_layer_' + str(i + 1) + '_' +
  448. str(j + 1)))
  449. self.residual_func_list.append(residual_func)
  450. elif j < i:
  451. pre_num_filters = in_channels[j]
  452. for k in range(i - j):
  453. if k == i - j - 1:
  454. residual_func = self.add_sublayer(
  455. "residual_{}_layer_{}_{}_{}".format(
  456. name, i + 1, j + 1, k + 1),
  457. ConvNormLayer(
  458. ch_in=pre_num_filters,
  459. ch_out=out_channels[i],
  460. filter_size=3,
  461. stride=2,
  462. norm_decay=norm_decay,
  463. freeze_norm=freeze_norm,
  464. act=None,
  465. name=name + '_layer_' + str(i + 1) + '_' +
  466. str(j + 1) + '_' + str(k + 1)))
  467. pre_num_filters = out_channels[i]
  468. else:
  469. residual_func = self.add_sublayer(
  470. "residual_{}_layer_{}_{}_{}".format(
  471. name, i + 1, j + 1, k + 1),
  472. ConvNormLayer(
  473. ch_in=pre_num_filters,
  474. ch_out=out_channels[j],
  475. filter_size=3,
  476. stride=2,
  477. norm_decay=norm_decay,
  478. freeze_norm=freeze_norm,
  479. act="relu",
  480. name=name + '_layer_' + str(i + 1) + '_' +
  481. str(j + 1) + '_' + str(k + 1)))
  482. pre_num_filters = out_channels[j]
  483. self.residual_func_list.append(residual_func)
  484. def forward(self, input):
  485. outs = []
  486. residual_func_idx = 0
  487. for i in range(self._actual_ch):
  488. residual = input[i]
  489. for j in range(len(self._in_channels)):
  490. if j > i:
  491. y = self.residual_func_list[residual_func_idx](input[j])
  492. residual_func_idx += 1
  493. y = F.interpolate(y, scale_factor=2**(j - i))
  494. residual = paddle.add(x=residual, y=y)
  495. elif j < i:
  496. y = input[j]
  497. for k in range(i - j):
  498. y = self.residual_func_list[residual_func_idx](y)
  499. residual_func_idx += 1
  500. residual = paddle.add(x=residual, y=y)
  501. residual = F.relu(residual)
  502. outs.append(residual)
  503. return outs
  504. @register
  505. class HRNet(nn.Layer):
  506. """
  507. HRNet, see https://arxiv.org/abs/1908.07919
  508. Args:
  509. width (int): the width of HRNet
  510. has_se (bool): whether to add SE block for each stage
  511. freeze_at (int): the stage to freeze
  512. freeze_norm (bool): whether to freeze norm in HRNet
  513. norm_decay (float): weight decay for normalization layer weights
  514. return_idx (List): the stage to return
  515. """
  516. def __init__(self,
  517. width=18,
  518. has_se=False,
  519. freeze_at=0,
  520. freeze_norm=True,
  521. norm_decay=0.,
  522. return_idx=[0, 1, 2, 3]):
  523. super(HRNet, self).__init__()
  524. self.width = width
  525. self.has_se = has_se
  526. if isinstance(return_idx, Integral):
  527. return_idx = [return_idx]
  528. assert len(return_idx) > 0, "need one or more return index"
  529. self.freeze_at = freeze_at
  530. self.return_idx = return_idx
  531. self.channels = {
  532. 18: [[18, 36], [18, 36, 72], [18, 36, 72, 144]],
  533. 30: [[30, 60], [30, 60, 120], [30, 60, 120, 240]],
  534. 32: [[32, 64], [32, 64, 128], [32, 64, 128, 256]],
  535. 40: [[40, 80], [40, 80, 160], [40, 80, 160, 320]],
  536. 44: [[44, 88], [44, 88, 176], [44, 88, 176, 352]],
  537. 48: [[48, 96], [48, 96, 192], [48, 96, 192, 384]],
  538. 60: [[60, 120], [60, 120, 240], [60, 120, 240, 480]],
  539. 64: [[64, 128], [64, 128, 256], [64, 128, 256, 512]]
  540. }
  541. channels_2, channels_3, channels_4 = self.channels[width]
  542. num_modules_2, num_modules_3, num_modules_4 = 1, 4, 3
  543. self._out_channels = channels_4
  544. self._out_strides = [4, 8, 16, 32]
  545. self.conv_layer1_1 = ConvNormLayer(
  546. ch_in=3,
  547. ch_out=64,
  548. filter_size=3,
  549. stride=2,
  550. norm_decay=norm_decay,
  551. freeze_norm=freeze_norm,
  552. act='relu',
  553. name="layer1_1")
  554. self.conv_layer1_2 = ConvNormLayer(
  555. ch_in=64,
  556. ch_out=64,
  557. filter_size=3,
  558. stride=2,
  559. norm_decay=norm_decay,
  560. freeze_norm=freeze_norm,
  561. act='relu',
  562. name="layer1_2")
  563. self.la1 = Layer1(
  564. num_channels=64,
  565. has_se=has_se,
  566. norm_decay=norm_decay,
  567. freeze_norm=freeze_norm,
  568. name="layer2")
  569. self.tr1 = TransitionLayer(
  570. in_channels=[256],
  571. out_channels=channels_2,
  572. norm_decay=norm_decay,
  573. freeze_norm=freeze_norm,
  574. name="tr1")
  575. self.st2 = Stage(
  576. num_channels=channels_2,
  577. num_modules=num_modules_2,
  578. num_filters=channels_2,
  579. has_se=self.has_se,
  580. norm_decay=norm_decay,
  581. freeze_norm=freeze_norm,
  582. name="st2")
  583. self.tr2 = TransitionLayer(
  584. in_channels=channels_2,
  585. out_channels=channels_3,
  586. norm_decay=norm_decay,
  587. freeze_norm=freeze_norm,
  588. name="tr2")
  589. self.st3 = Stage(
  590. num_channels=channels_3,
  591. num_modules=num_modules_3,
  592. num_filters=channels_3,
  593. has_se=self.has_se,
  594. norm_decay=norm_decay,
  595. freeze_norm=freeze_norm,
  596. name="st3")
  597. self.tr3 = TransitionLayer(
  598. in_channels=channels_3,
  599. out_channels=channels_4,
  600. norm_decay=norm_decay,
  601. freeze_norm=freeze_norm,
  602. name="tr3")
  603. self.st4 = Stage(
  604. num_channels=channels_4,
  605. num_modules=num_modules_4,
  606. num_filters=channels_4,
  607. has_se=self.has_se,
  608. norm_decay=norm_decay,
  609. freeze_norm=freeze_norm,
  610. multi_scale_output=len(return_idx) > 1,
  611. name="st4")
  612. def forward(self, inputs):
  613. x = inputs['image']
  614. conv1 = self.conv_layer1_1(x)
  615. conv2 = self.conv_layer1_2(conv1)
  616. la1 = self.la1(conv2)
  617. tr1 = self.tr1([la1])
  618. st2 = self.st2(tr1)
  619. tr2 = self.tr2(st2)
  620. st3 = self.st3(tr2)
  621. tr3 = self.tr3(st3)
  622. st4 = self.st4(tr3)
  623. res = []
  624. for i, layer in enumerate(st4):
  625. if i == self.freeze_at:
  626. layer.stop_gradient = True
  627. if i in self.return_idx:
  628. res.append(layer)
  629. return res
  630. @property
  631. def out_shape(self):
  632. return [
  633. ShapeSpec(
  634. channels=self._out_channels[i], stride=self._out_strides[i])
  635. for i in self.return_idx
  636. ]