提交 9b61df62 编写于 作者: W weishengyu

replace add_sublayer with nn.Sequential

上级 a16de6da
...@@ -54,8 +54,7 @@ class ConvBNLayer(TheseusLayer): ...@@ -54,8 +54,7 @@ class ConvBNLayer(TheseusLayer):
filter_size, filter_size,
stride=1, stride=1,
groups=1, groups=1,
act="relu", act="relu"):
name=None):
super(ConvBNLayer, self).__init__() super(ConvBNLayer, self).__init__()
self._conv = nn.Conv2D( self._conv = nn.Conv2D(
...@@ -66,7 +65,6 @@ class ConvBNLayer(TheseusLayer): ...@@ -66,7 +65,6 @@ class ConvBNLayer(TheseusLayer):
padding=(filter_size - 1) // 2, padding=(filter_size - 1) // 2,
groups=groups, groups=groups,
bias_attr=False) bias_attr=False)
bn_name = name + '_bn'
self._batch_norm = nn.BatchNorm( self._batch_norm = nn.BatchNorm(
num_filters, num_filters,
act=act) act=act)
...@@ -81,19 +79,16 @@ class Layer1(TheseusLayer): ...@@ -81,19 +79,16 @@ class Layer1(TheseusLayer):
def __init__(self, num_channels, has_se=False, name=None): def __init__(self, num_channels, has_se=False, name=None):
super(Layer1, self).__init__() super(Layer1, self).__init__()
self.bottleneck_block_list = [] self.bottleneck_block_list = nn.Sequential(*[
BottleneckBlock(
for i in range(4): num_channels=num_channels if i == 0 else 256,
bottleneck_block = self.add_sublayer( num_filters=64,
"bb_{}_{}".format(name, i + 1), has_se=has_se,
BottleneckBlock( stride=1,
num_channels=num_channels if i == 0 else 256, downsample=True if i == 0 else False,
num_filters=64, name=name + '_' + str(i + 1))
has_se=has_se, for i in range(4)
stride=1, ])
downsample=True if i == 0 else False,
name=name + '_' + str(i + 1)))
self.bottleneck_block_list.append(bottleneck_block)
def forward(self, x, res_dict=None): def forward(self, x, res_dict=None):
y = x y = x
......
Markdown is supported
0% .
You are about to add 0 people to the discussion. Proceed with caution.
先完成此消息的编辑!
想要评论请 注册