572 lines
17 KiB
Python
572 lines
17 KiB
Python
# copyright (c) 2021 PaddlePaddle Authors. All Rights Reserve.
|
|
#
|
|
# 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.
|
|
from __future__ import annotations
|
|
|
|
import math
|
|
from typing import TYPE_CHECKING
|
|
|
|
from typing_extensions import Unpack
|
|
|
|
import paddle
|
|
from paddle import nn
|
|
from paddle.base.param_attr import ParamAttr
|
|
from paddle.nn import (
|
|
AdaptiveAvgPool2D,
|
|
AvgPool2D,
|
|
BatchNorm,
|
|
Conv2D,
|
|
Dropout,
|
|
Linear,
|
|
MaxPool2D,
|
|
)
|
|
from paddle.nn.initializer import Uniform
|
|
from paddle.utils.download import get_weights_path_from_url
|
|
|
|
if TYPE_CHECKING:
|
|
from typing import Literal, TypedDict
|
|
|
|
from typing_extensions import NotRequired
|
|
|
|
from paddle import Tensor
|
|
from paddle._typing import Size2
|
|
|
|
_DenseNetArch = Literal[
|
|
"densenet121",
|
|
"densenet161",
|
|
"densenet169",
|
|
"densenet201",
|
|
"densenet264",
|
|
]
|
|
|
|
class _DenseNetOptions(TypedDict):
|
|
bn_size: NotRequired[int]
|
|
dropout: NotRequired[float]
|
|
num_classes: NotRequired[int]
|
|
with_pool: NotRequired[bool]
|
|
|
|
|
|
__all__ = []
|
|
|
|
model_urls: dict[str, tuple[str, str]] = {
|
|
'densenet121': (
|
|
'https://paddle-imagenet-models-name.bj.bcebos.com/dygraph/DenseNet121_pretrained.pdparams',
|
|
'db1b239ed80a905290fd8b01d3af08e4',
|
|
),
|
|
'densenet161': (
|
|
'https://paddle-imagenet-models-name.bj.bcebos.com/dygraph/DenseNet161_pretrained.pdparams',
|
|
'62158869cb315098bd25ddbfd308a853',
|
|
),
|
|
'densenet169': (
|
|
'https://paddle-imagenet-models-name.bj.bcebos.com/dygraph/DenseNet169_pretrained.pdparams',
|
|
'82cc7c635c3f19098c748850efb2d796',
|
|
),
|
|
'densenet201': (
|
|
'https://paddle-imagenet-models-name.bj.bcebos.com/dygraph/DenseNet201_pretrained.pdparams',
|
|
'16ca29565a7712329cf9e36e02caaf58',
|
|
),
|
|
'densenet264': (
|
|
'https://paddle-imagenet-models-name.bj.bcebos.com/dygraph/DenseNet264_pretrained.pdparams',
|
|
'3270ce516b85370bba88cfdd9f60bff4',
|
|
),
|
|
}
|
|
|
|
|
|
class BNACConvLayer(nn.Layer):
|
|
def __init__(
|
|
self,
|
|
num_channels: int,
|
|
num_filters: int,
|
|
filter_size: Size2,
|
|
stride: Size2 = 1,
|
|
pad: Size2 = 0,
|
|
groups: int = 1,
|
|
act: str = "relu",
|
|
) -> None:
|
|
super().__init__()
|
|
self._batch_norm = BatchNorm(num_channels, act=act)
|
|
|
|
self._conv = Conv2D(
|
|
in_channels=num_channels,
|
|
out_channels=num_filters,
|
|
kernel_size=filter_size,
|
|
stride=stride,
|
|
padding=pad,
|
|
groups=groups,
|
|
weight_attr=ParamAttr(),
|
|
bias_attr=False,
|
|
)
|
|
|
|
def forward(self, input: Tensor) -> Tensor:
|
|
y = self._batch_norm(input)
|
|
y = self._conv(y)
|
|
return y
|
|
|
|
|
|
class DenseLayer(nn.Layer):
|
|
dropout: float
|
|
|
|
def __init__(
|
|
self, num_channels: int, growth_rate: int, bn_size: int, dropout: float
|
|
) -> None:
|
|
super().__init__()
|
|
self.dropout = dropout
|
|
|
|
self.bn_ac_func1 = BNACConvLayer(
|
|
num_channels=num_channels,
|
|
num_filters=bn_size * growth_rate,
|
|
filter_size=1,
|
|
pad=0,
|
|
stride=1,
|
|
)
|
|
|
|
self.bn_ac_func2 = BNACConvLayer(
|
|
num_channels=bn_size * growth_rate,
|
|
num_filters=growth_rate,
|
|
filter_size=3,
|
|
pad=1,
|
|
stride=1,
|
|
)
|
|
|
|
if dropout:
|
|
self.dropout_func = Dropout(p=dropout, mode="downscale_in_infer")
|
|
|
|
def forward(self, input: Tensor) -> Tensor:
|
|
conv = self.bn_ac_func1(input)
|
|
conv = self.bn_ac_func2(conv)
|
|
if self.dropout:
|
|
conv = self.dropout_func(conv)
|
|
conv = paddle.concat([input, conv], axis=1)
|
|
return conv
|
|
|
|
|
|
class DenseBlock(nn.Layer):
|
|
dropout: float
|
|
|
|
def __init__(
|
|
self,
|
|
num_channels: int,
|
|
num_layers: int,
|
|
bn_size: int,
|
|
growth_rate: int,
|
|
dropout: float,
|
|
name: str | None = None,
|
|
) -> None:
|
|
super().__init__()
|
|
self.dropout = dropout
|
|
self.dense_layer_func = []
|
|
|
|
pre_channel = num_channels
|
|
for layer in range(num_layers):
|
|
self.dense_layer_func.append(
|
|
self.add_sublayer(
|
|
f"{name}_{layer + 1}",
|
|
DenseLayer(
|
|
num_channels=pre_channel,
|
|
growth_rate=growth_rate,
|
|
bn_size=bn_size,
|
|
dropout=dropout,
|
|
),
|
|
)
|
|
)
|
|
pre_channel = pre_channel + growth_rate
|
|
|
|
def forward(self, input: Tensor) -> Tensor:
|
|
conv = input
|
|
for func in self.dense_layer_func:
|
|
conv = func(conv)
|
|
return conv
|
|
|
|
|
|
class TransitionLayer(nn.Layer):
|
|
def __init__(self, num_channels: int, num_output_features: int) -> None:
|
|
super().__init__()
|
|
|
|
self.conv_ac_func = BNACConvLayer(
|
|
num_channels=num_channels,
|
|
num_filters=num_output_features,
|
|
filter_size=1,
|
|
pad=0,
|
|
stride=1,
|
|
)
|
|
|
|
self.pool2d_avg = AvgPool2D(kernel_size=2, stride=2, padding=0)
|
|
|
|
def forward(self, input: Tensor) -> Tensor:
|
|
y = self.conv_ac_func(input)
|
|
y = self.pool2d_avg(y)
|
|
return y
|
|
|
|
|
|
class ConvBNLayer(nn.Layer):
|
|
def __init__(
|
|
self,
|
|
num_channels: int,
|
|
num_filters: int,
|
|
filter_size: Size2,
|
|
stride: Size2 = 1,
|
|
pad: Size2 = 0,
|
|
groups: int = 1,
|
|
act: str = "relu",
|
|
) -> None:
|
|
super().__init__()
|
|
|
|
self._conv = Conv2D(
|
|
in_channels=num_channels,
|
|
out_channels=num_filters,
|
|
kernel_size=filter_size,
|
|
stride=stride,
|
|
padding=pad,
|
|
groups=groups,
|
|
weight_attr=ParamAttr(),
|
|
bias_attr=False,
|
|
)
|
|
self._batch_norm = BatchNorm(num_filters, act=act)
|
|
|
|
def forward(self, input: Tensor) -> Tensor:
|
|
y = self._conv(input)
|
|
y = self._batch_norm(y)
|
|
return y
|
|
|
|
|
|
class DenseNet(nn.Layer):
|
|
"""DenseNet model from
|
|
`"Densely Connected Convolutional Networks" <https://arxiv.org/pdf/1608.06993.pdf>`_.
|
|
|
|
Args:
|
|
layers (int, optional): Layers of DenseNet. Default: 121.
|
|
bn_size (int, optional): Expansion of growth rate in the middle layer. Default: 4.
|
|
dropout (float, optional): Dropout rate. Default: :math:`0.0`.
|
|
num_classes (int, optional): Output dim of last fc layer. If num_classes <= 0, last fc layer
|
|
will not be defined. Default: 1000.
|
|
with_pool (bool, optional): Use pool before the last fc layer or not. Default: True.
|
|
|
|
Returns:
|
|
:ref:`api_paddle_nn_Layer`. An instance of DenseNet model.
|
|
|
|
Examples:
|
|
.. code-block:: pycon
|
|
|
|
>>> import paddle
|
|
>>> from paddle.vision.models import DenseNet
|
|
|
|
>>> # Build model
|
|
>>> densenet = DenseNet()
|
|
|
|
>>> x = paddle.rand([1, 3, 224, 224])
|
|
>>> out = densenet(x)
|
|
|
|
>>> print(out.shape)
|
|
paddle.Size([1, 1000])
|
|
"""
|
|
|
|
num_classes: int
|
|
with_pool: bool
|
|
|
|
def __init__(
|
|
self,
|
|
layers: int = 121,
|
|
bn_size: int = 4,
|
|
dropout: float = 0.0,
|
|
num_classes: int = 1000,
|
|
with_pool: bool = True,
|
|
) -> None:
|
|
super().__init__()
|
|
self.num_classes = num_classes
|
|
self.with_pool = with_pool
|
|
supported_layers = [121, 161, 169, 201, 264]
|
|
assert layers in supported_layers, (
|
|
f"supported layers are {supported_layers} but input layer is {layers}"
|
|
)
|
|
densenet_spec = {
|
|
121: (64, 32, [6, 12, 24, 16]),
|
|
161: (96, 48, [6, 12, 36, 24]),
|
|
169: (64, 32, [6, 12, 32, 32]),
|
|
201: (64, 32, [6, 12, 48, 32]),
|
|
264: (64, 32, [6, 12, 64, 48]),
|
|
}
|
|
num_init_features, growth_rate, block_config = densenet_spec[layers]
|
|
|
|
self.conv1_func = ConvBNLayer(
|
|
num_channels=3,
|
|
num_filters=num_init_features,
|
|
filter_size=7,
|
|
stride=2,
|
|
pad=3,
|
|
act='relu',
|
|
)
|
|
self.pool2d_max = MaxPool2D(kernel_size=3, stride=2, padding=1)
|
|
self.block_config = block_config
|
|
self.dense_block_func_list = []
|
|
self.transition_func_list = []
|
|
pre_num_channels = num_init_features
|
|
num_features = num_init_features
|
|
for i, num_layers in enumerate(block_config):
|
|
self.dense_block_func_list.append(
|
|
self.add_sublayer(
|
|
f"db_conv_{i + 2}",
|
|
DenseBlock(
|
|
num_channels=pre_num_channels,
|
|
num_layers=num_layers,
|
|
bn_size=bn_size,
|
|
growth_rate=growth_rate,
|
|
dropout=dropout,
|
|
name='conv' + str(i + 2),
|
|
),
|
|
)
|
|
)
|
|
|
|
num_features = num_features + num_layers * growth_rate
|
|
pre_num_channels = num_features
|
|
|
|
if i != len(block_config) - 1:
|
|
self.transition_func_list.append(
|
|
self.add_sublayer(
|
|
f"tr_conv{i + 2}_blk",
|
|
TransitionLayer(
|
|
num_channels=pre_num_channels,
|
|
num_output_features=num_features // 2,
|
|
),
|
|
)
|
|
)
|
|
pre_num_channels = num_features // 2
|
|
num_features = num_features // 2
|
|
|
|
self.batch_norm = BatchNorm(num_features, act="relu")
|
|
if self.with_pool:
|
|
self.pool2d_avg = AdaptiveAvgPool2D(1)
|
|
|
|
if self.num_classes > 0:
|
|
stdv = 1.0 / math.sqrt(num_features * 1.0)
|
|
self.out = Linear(
|
|
num_features,
|
|
num_classes,
|
|
weight_attr=ParamAttr(initializer=Uniform(-stdv, stdv)),
|
|
bias_attr=ParamAttr(),
|
|
)
|
|
|
|
def forward(self, input: Tensor) -> Tensor:
|
|
conv = self.conv1_func(input)
|
|
conv = self.pool2d_max(conv)
|
|
|
|
for i, num_layers in enumerate(self.block_config):
|
|
conv = self.dense_block_func_list[i](conv)
|
|
if i != len(self.block_config) - 1:
|
|
conv = self.transition_func_list[i](conv)
|
|
|
|
conv = self.batch_norm(conv)
|
|
|
|
if self.with_pool:
|
|
y = self.pool2d_avg(conv)
|
|
|
|
if self.num_classes > 0:
|
|
y = paddle.flatten(y, start_axis=1, stop_axis=-1)
|
|
y = self.out(y)
|
|
|
|
return y
|
|
|
|
|
|
def _densenet(
|
|
arch: _DenseNetArch,
|
|
layers: int,
|
|
pretrained: bool,
|
|
**kwargs: Unpack[_DenseNetOptions],
|
|
) -> DenseNet:
|
|
model = DenseNet(layers=layers, **kwargs)
|
|
if pretrained:
|
|
assert arch in model_urls, (
|
|
f"{arch} model do not have a pretrained model now, you should set pretrained=False"
|
|
)
|
|
weight_path = get_weights_path_from_url(
|
|
model_urls[arch][0], model_urls[arch][1]
|
|
)
|
|
|
|
param = paddle.load(weight_path)
|
|
model.set_dict(param)
|
|
|
|
return model
|
|
|
|
|
|
def densenet121(
|
|
pretrained: bool = False, **kwargs: Unpack[_DenseNetOptions]
|
|
) -> DenseNet:
|
|
"""DenseNet 121-layer model from
|
|
`"Densely Connected Convolutional Networks" <https://arxiv.org/pdf/1608.06993.pdf>`_.
|
|
|
|
Args:
|
|
pretrained (bool, optional): Whether to load pre-trained weights. If True, returns a model pre-trained
|
|
on ImageNet. Default: False.
|
|
**kwargs (optional): Additional keyword arguments. For details, please refer to :ref:`DenseNet <api_paddle_vision_models_DenseNet>`.
|
|
|
|
Returns:
|
|
:ref:`api_paddle_nn_Layer`. An instance of DenseNet 121-layer model.
|
|
|
|
Examples:
|
|
.. code-block:: pycon
|
|
|
|
>>> import paddle
|
|
>>> from paddle.vision.models import densenet121
|
|
|
|
>>> # Build model
|
|
>>> model = densenet121()
|
|
|
|
>>> # Build model and load imagenet pretrained weight
|
|
>>> # model = densenet121(pretrained=True)
|
|
|
|
>>> x = paddle.rand([1, 3, 224, 224])
|
|
>>> out = model(x)
|
|
|
|
>>> print(out.shape)
|
|
paddle.Size([1, 1000])
|
|
"""
|
|
return _densenet('densenet121', 121, pretrained, **kwargs)
|
|
|
|
|
|
def densenet161(
|
|
pretrained: bool = False, **kwargs: Unpack[_DenseNetOptions]
|
|
) -> DenseNet:
|
|
"""DenseNet 161-layer model from
|
|
`"Densely Connected Convolutional Networks" <https://arxiv.org/pdf/1608.06993.pdf>`_.
|
|
|
|
Args:
|
|
pretrained (bool, optional): Whether to load pre-trained weights. If True, returns a model pre-trained
|
|
on ImageNet. Default: False.
|
|
**kwargs (optional): Additional keyword arguments. For details, please refer to :ref:`DenseNet <api_paddle_vision_models_DenseNet>`.
|
|
|
|
Returns:
|
|
:ref:`api_paddle_nn_Layer`. An instance of DenseNet 161-layer model.
|
|
|
|
Examples:
|
|
.. code-block:: pycon
|
|
|
|
>>> import paddle
|
|
>>> from paddle.vision.models import densenet161
|
|
|
|
>>> # Build model
|
|
>>> model = densenet161()
|
|
|
|
>>> # Build model and load imagenet pretrained weight
|
|
>>> # model = densenet161(pretrained=True)
|
|
|
|
>>> x = paddle.rand([1, 3, 224, 224])
|
|
>>> out = model(x)
|
|
|
|
>>> print(out.shape)
|
|
paddle.Size([1, 1000])
|
|
"""
|
|
return _densenet('densenet161', 161, pretrained, **kwargs)
|
|
|
|
|
|
def densenet169(
|
|
pretrained: bool = False, **kwargs: Unpack[_DenseNetOptions]
|
|
) -> DenseNet:
|
|
"""DenseNet 169-layer model from
|
|
`"Densely Connected Convolutional Networks" <https://arxiv.org/pdf/1608.06993.pdf>`_.
|
|
|
|
Args:
|
|
pretrained (bool, optional): Whether to load pre-trained weights. If True, returns a model pre-trained
|
|
on ImageNet. Default: False.
|
|
**kwargs (optional): Additional keyword arguments. For details, please refer to :ref:`DenseNet <api_paddle_vision_models_DenseNet>`.
|
|
|
|
Returns:
|
|
:ref:`api_paddle_nn_Layer`. An instance of DenseNet 169-layer model.
|
|
|
|
Examples:
|
|
.. code-block:: pycon
|
|
|
|
>>> import paddle
|
|
>>> from paddle.vision.models import densenet169
|
|
|
|
>>> # Build model
|
|
>>> model = densenet169()
|
|
|
|
>>> # Build model and load imagenet pretrained weight
|
|
>>> # model = densenet169(pretrained=True)
|
|
|
|
>>> x = paddle.rand([1, 3, 224, 224])
|
|
>>> out = model(x)
|
|
|
|
>>> print(out.shape)
|
|
paddle.Size([1, 1000])
|
|
"""
|
|
return _densenet('densenet169', 169, pretrained, **kwargs)
|
|
|
|
|
|
def densenet201(
|
|
pretrained: bool = False, **kwargs: Unpack[_DenseNetOptions]
|
|
) -> DenseNet:
|
|
"""DenseNet 201-layer model from
|
|
`"Densely Connected Convolutional Networks" <https://arxiv.org/pdf/1608.06993.pdf>`_.
|
|
|
|
Args:
|
|
pretrained (bool, optional): Whether to load pre-trained weights. If True, returns a model pre-trained
|
|
on ImageNet. Default: False.
|
|
**kwargs (optional): Additional keyword arguments. For details, please refer to :ref:`DenseNet <api_paddle_vision_models_DenseNet>`.
|
|
|
|
Returns:
|
|
:ref:`api_paddle_nn_Layer`. An instance of DenseNet 201-layer model.
|
|
|
|
Examples:
|
|
.. code-block:: pycon
|
|
|
|
>>> import paddle
|
|
>>> from paddle.vision.models import densenet201
|
|
|
|
>>> # Build model
|
|
>>> model = densenet201()
|
|
|
|
>>> # Build model and load imagenet pretrained weight
|
|
>>> # model = densenet201(pretrained=True)
|
|
>>> x = paddle.rand([1, 3, 224, 224])
|
|
>>> out = model(x)
|
|
|
|
>>> print(out.shape)
|
|
paddle.Size([1, 1000])
|
|
"""
|
|
return _densenet('densenet201', 201, pretrained, **kwargs)
|
|
|
|
|
|
def densenet264(
|
|
pretrained: bool = False, **kwargs: Unpack[_DenseNetOptions]
|
|
) -> DenseNet:
|
|
"""DenseNet 264-layer model from
|
|
`"Densely Connected Convolutional Networks" <https://arxiv.org/pdf/1608.06993.pdf>`_.
|
|
|
|
Args:
|
|
pretrained (bool, optional): Whether to load pre-trained weights. If True, returns a model pre-trained
|
|
on ImageNet. Default: False.
|
|
**kwargs (optional): Additional keyword arguments. For details, please refer to :ref:`DenseNet <api_paddle_vision_models_DenseNet>`.
|
|
|
|
Returns:
|
|
:ref:`api_paddle_nn_Layer`. An instance of DenseNet 264-layer model.
|
|
|
|
Examples:
|
|
.. code-block:: pycon
|
|
|
|
>>> import paddle
|
|
>>> from paddle.vision.models import densenet264
|
|
|
|
>>> # Build model
|
|
>>> model = densenet264()
|
|
|
|
>>> # Build model and load imagenet pretrained weight
|
|
>>> # model = densenet264(pretrained=True)
|
|
|
|
>>> x = paddle.rand([1, 3, 224, 224])
|
|
>>> out = model(x)
|
|
|
|
>>> print(out.shape)
|
|
paddle.Size([1, 1000])
|
|
"""
|
|
return _densenet('densenet264', 264, pretrained, **kwargs)
|