Spaces:
Configuration error
Configuration error
| # Copyright (c) 2021 PaddlePaddle Authors. All Rights Reserved. | |
| # | |
| # Licensed under the Apache License, Version 2.0 (the "License"); | |
| # you may not use this file except in compliance with the License. | |
| # You may obtain a copy of the License at | |
| # | |
| # http://www.apache.org/licenses/LICENSE-2.0 | |
| # | |
| # Unless required by applicable law or agreed to in writing, software | |
| # distributed under the License is distributed on an "AS IS" BASIS, | |
| # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. | |
| # See the License for the specific language governing permissions and | |
| # limitations under the License. | |
| import math | |
| import paddle | |
| import paddle.nn as nn | |
| from paddleseg.utils import utils | |
| from paddleseg.cvlibs import manager, param_init | |
| from paddleseg.models.layers.layer_libs import SyncBatchNorm | |
| __all__ = ["STDC1", "STDC2"] | |
| class STDCNet(nn.Layer): | |
| """ | |
| The STDCNet implementation based on PaddlePaddle. | |
| The original article refers to Meituan | |
| Fan, Mingyuan, et al. "Rethinking BiSeNet For Real-time Semantic Segmentation." | |
| (https://arxiv.org/abs/2104.13188) | |
| Args: | |
| base(int, optional): base channels. Default: 64. | |
| layers(list, optional): layers numbers list. It determines STDC block numbers of STDCNet's stage3\4\5. Defualt: [4, 5, 3]. | |
| block_num(int,optional): block_num of features block. Default: 4. | |
| type(str,optional): feature fusion method "cat"/"add". Default: "cat". | |
| relative_lr(float,optional): parameters here receive a different learning rate when updating. The effective | |
| learning rate is the prodcut of relative_lr and the global learning rate. Default: 1.0. | |
| pretrained(str, optional): the path of pretrained model. | |
| """ | |
| def __init__(self, | |
| base=64, | |
| layers=[4, 5, 3], | |
| block_num=4, | |
| type="cat", | |
| relative_lr=1.0, | |
| pretrained=None): | |
| super(STDCNet, self).__init__() | |
| if type == "cat": | |
| block = CatBottleneck | |
| elif type == "add": | |
| block = AddBottleneck | |
| self.layers = layers | |
| self.feat_channels = [base // 2, base, base * 4, base * 8, base * 16] | |
| self.features = self._make_layers(base, layers, block_num, block, relative_lr) | |
| self.pretrained = pretrained | |
| self.init_weight() | |
| def forward(self, x): | |
| """ | |
| forward function for feature extract. | |
| """ | |
| out_feats = [] | |
| x = self.features[0](x) | |
| out_feats.append(x) | |
| x = self.features[1](x) | |
| out_feats.append(x) | |
| idx = [[2, 2 + self.layers[0]], | |
| [2 + self.layers[0], 2 + sum(self.layers[0:2])], | |
| [2 + sum(self.layers[0:2]), 2 + sum(self.layers)]] | |
| for start_idx, end_idx in idx: | |
| for i in range(start_idx, end_idx): | |
| x = self.features[i](x) | |
| out_feats.append(x) | |
| return out_feats | |
| def _make_layers(self, base, layers, block_num, block, relative_lr): | |
| features = [] | |
| features += [ConvBNRelu(3, base // 2, 3, 2, relative_lr)] | |
| features += [ConvBNRelu(base // 2, base, 3, 2, relative_lr)] | |
| for i, layer in enumerate(layers): | |
| for j in range(layer): | |
| if i == 0 and j == 0: | |
| features.append(block(base, base * 4, block_num, 2, relative_lr)) | |
| elif j == 0: | |
| features.append( | |
| block(base * int(math.pow(2, i + 1)), base * int( | |
| math.pow(2, i + 2)), block_num, 2, relative_lr)) | |
| else: | |
| features.append( | |
| block(base * int(math.pow(2, i + 2)), base * int( | |
| math.pow(2, i + 2)), block_num, 1, relative_lr)) | |
| return nn.Sequential(*features) | |
| def init_weight(self): | |
| for layer in self.sublayers(): | |
| if isinstance(layer, nn.Conv2D): | |
| param_init.normal_init(layer.weight, std=0.001) | |
| elif isinstance(layer, (nn.BatchNorm, nn.SyncBatchNorm)): | |
| param_init.constant_init(layer.weight, value=1.0) | |
| param_init.constant_init(layer.bias, value=0.0) | |
| if self.pretrained is not None: | |
| utils.load_pretrained_model(self, self.pretrained) | |
| class ConvBNRelu(nn.Layer): | |
| def __init__(self, in_planes, out_planes, kernel=3, stride=1, relative_lr=1.0): | |
| super(ConvBNRelu, self).__init__() | |
| param_attr = paddle.ParamAttr(learning_rate=relative_lr) | |
| self.conv = nn.Conv2D( | |
| in_planes, | |
| out_planes, | |
| kernel_size=kernel, | |
| stride=stride, | |
| padding=kernel // 2, | |
| weight_attr=param_attr, | |
| bias_attr=False) | |
| self.bn = nn.BatchNorm2D( | |
| out_planes, | |
| weight_attr=param_attr, | |
| bias_attr=param_attr | |
| ) | |
| self.relu = nn.ReLU() | |
| def forward(self, x): | |
| out = self.relu(self.bn(self.conv(x))) | |
| return out | |
| class AddBottleneck(nn.Layer): | |
| def __init__(self, in_planes, out_planes, block_num=3, stride=1, relative_lr=1.0): | |
| super(AddBottleneck, self).__init__() | |
| assert block_num > 1, "block number should be larger than 1." | |
| self.conv_list = nn.LayerList() | |
| self.stride = stride | |
| param_attr = paddle.ParamAttr(learning_rate=relative_lr) | |
| if stride == 2: | |
| self.avd_layer = nn.Sequential( | |
| nn.Conv2D( | |
| out_planes // 2, | |
| out_planes // 2, | |
| kernel_size=3, | |
| stride=2, | |
| padding=1, | |
| groups=out_planes // 2, | |
| weight_attr=param_attr, | |
| bias_attr=False), | |
| nn.BatchNorm2D(out_planes // 2, weight_attr=param_attr, bias_attr=param_attr), ) | |
| self.skip = nn.Sequential( | |
| nn.Conv2D( | |
| in_planes, | |
| in_planes, | |
| kernel_size=3, | |
| stride=2, | |
| padding=1, | |
| groups=in_planes, | |
| weight_attr=param_attr, | |
| bias_attr=False), | |
| nn.BatchNorm2D(in_planes, weight_attr=param_attr, bias_attr=param_attr), | |
| nn.Conv2D( | |
| in_planes, out_planes, kernel_size=1, bias_attr=False, | |
| weight_attr=param_attr | |
| ), | |
| nn.BatchNorm2D(out_planes, weight_attr=param_attr, bias_attr=param_attr), ) | |
| stride = 1 | |
| for idx in range(block_num): | |
| if idx == 0: | |
| self.conv_list.append( | |
| ConvBNRelu( | |
| in_planes, out_planes // 2, kernel=1, relative_lr=relative_lr)) | |
| elif idx == 1 and block_num == 2: | |
| self.conv_list.append( | |
| ConvBNRelu( | |
| out_planes // 2, out_planes // 2, stride=stride, relative_lr=relative_lr)) | |
| elif idx == 1 and block_num > 2: | |
| self.conv_list.append( | |
| ConvBNRelu( | |
| out_planes // 2, out_planes // 4, stride=stride, relative_lr=relative_lr)) | |
| elif idx < block_num - 1: | |
| self.conv_list.append( | |
| ConvBNRelu(out_planes // int(math.pow(2, idx)), out_planes | |
| // int(math.pow(2, idx + 1)), relative_lr=relative_lr) | |
| ) | |
| else: | |
| self.conv_list.append( | |
| ConvBNRelu(out_planes // int(math.pow(2, idx)), out_planes | |
| // int(math.pow(2, idx))), relative_lr=relative_lr | |
| ) | |
| def forward(self, x): | |
| out_list = [] | |
| out = x | |
| for idx, conv in enumerate(self.conv_list): | |
| if idx == 0 and self.stride == 2: | |
| out = self.avd_layer(conv(out)) | |
| else: | |
| out = conv(out) | |
| out_list.append(out) | |
| if self.stride == 2: | |
| x = self.skip(x) | |
| return paddle.concat(out_list, axis=1) + x | |
| class CatBottleneck(nn.Layer): | |
| def __init__(self, in_planes, out_planes, block_num=3, stride=1, relative_lr=1.0): | |
| super(CatBottleneck, self).__init__() | |
| assert block_num > 1, "block number should be larger than 1." | |
| self.conv_list = nn.LayerList() | |
| self.stride = stride | |
| param_attr = paddle.ParamAttr(learning_rate=relative_lr) | |
| if stride == 2: | |
| self.avd_layer = nn.Sequential( | |
| nn.Conv2D( | |
| out_planes // 2, | |
| out_planes // 2, | |
| kernel_size=3, | |
| stride=2, | |
| padding=1, | |
| groups=out_planes // 2, | |
| weight_attr=param_attr, | |
| bias_attr=False), | |
| nn.BatchNorm2D(out_planes // 2, weight_attr=param_attr, bias_attr=param_attr), ) | |
| self.skip = nn.AvgPool2D(kernel_size=3, stride=2, padding=1) | |
| stride = 1 | |
| for idx in range(block_num): | |
| if idx == 0: | |
| self.conv_list.append( | |
| ConvBNRelu( | |
| in_planes, out_planes // 2, kernel=1, relative_lr=relative_lr)) | |
| elif idx == 1 and block_num == 2: | |
| self.conv_list.append( | |
| ConvBNRelu( | |
| out_planes // 2, out_planes // 2, stride=stride, relative_lr=relative_lr)) | |
| elif idx == 1 and block_num > 2: | |
| self.conv_list.append( | |
| ConvBNRelu( | |
| out_planes // 2, out_planes // 4, stride=stride, relative_lr=relative_lr)) | |
| elif idx < block_num - 1: | |
| self.conv_list.append( | |
| ConvBNRelu(out_planes // int(math.pow(2, idx)), out_planes | |
| // int(math.pow(2, idx + 1)), relative_lr=relative_lr)) | |
| else: | |
| self.conv_list.append( | |
| ConvBNRelu(out_planes // int(math.pow(2, idx)), out_planes | |
| // int(math.pow(2, idx)), relative_lr=relative_lr)) | |
| def forward(self, x): | |
| out_list = [] | |
| out1 = self.conv_list[0](x) | |
| for idx, conv in enumerate(self.conv_list[1:]): | |
| if idx == 0: | |
| if self.stride == 2: | |
| out = conv(self.avd_layer(out1)) | |
| else: | |
| out = conv(out1) | |
| else: | |
| out = conv(out) | |
| out_list.append(out) | |
| if self.stride == 2: | |
| out1 = self.skip(out1) | |
| out_list.insert(0, out1) | |
| out = paddle.concat(out_list, axis=1) | |
| return out | |
| def STDC2(**kwargs): | |
| model = STDCNet(base=64, layers=[4, 5, 3], **kwargs) | |
| return model | |
| def STDC1(**kwargs): | |
| model = STDCNet(base=64, layers=[2, 2, 2], **kwargs) | |
| return model | |