# Copyright (c) 2020 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.
from __future__ import absolute_import
from __future__ import division
from __future__ import print_function
from paddle import fluid
from paddle.fluid.param_attr import ParamAttr
from paddle.fluid.initializer import Uniform
from ppdet.core.workspace import register
__all__ = ['Hourglass']
def kaiming_init(input, filter_size):
fan_in = input.shape[1]
std = (1.0 / (fan_in * filter_size * filter_size))**0.5
return Uniform(0. - std, std)
def _conv_norm(x,
k,
out_dim,
stride=1,
pad=0,
groups=None,
with_bn=True,
bn_act=None,
ind=None,
name=None):
conv_name = "_conv" if ind is None else "_conv" + str(ind)
bn_name = "_bn" if ind is None else "_bn" + str(ind)
conv = fluid.layers.conv2d(
input=x,
filter_size=k,
num_filters=out_dim,
stride=stride,
padding=pad,
groups=groups,
param_attr=ParamAttr(
name=name + conv_name + "_weight", initializer=kaiming_init(x, k)),
bias_attr=ParamAttr(
name=name + conv_name + "_bias", initializer=kaiming_init(x, k))
if not with_bn else False,
name=name + '_output')
if with_bn:
pattr = ParamAttr(name=name + bn_name + '_weight')
battr = ParamAttr(name=name + bn_name + '_bias')
out = fluid.layers.batch_norm(
input=conv,
act=bn_act,
name=name + '_bn_output',
param_attr=pattr,
bias_attr=battr,
moving_mean_name=name + bn_name + '_running_mean',
moving_variance_name=name + bn_name +
'_running_var') if with_bn else conv
else:
out = fluid.layers.relu(conv)
return out
def residual_block(x, out_dim, k=3, stride=1, name=None):
p = (k - 1) // 2
conv1 = _conv_norm(
x, k, out_dim, pad=p, stride=stride, bn_act='relu', ind=1, name=name)
conv2 = _conv_norm(conv1, k, out_dim, pad=p, ind=2, name=name)
skip = _conv_norm(
x, 1, out_dim, stride=stride,
name=name + '_skip') if stride != 1 or x.shape[1] != out_dim else x
return fluid.layers.elementwise_add(
x=skip, y=conv2, act='relu', name=name + "_add")
def fire_block(x, out_dim, sr=2, stride=1, name=None):
conv1 = _conv_norm(x, 1, out_dim // sr, ind=1, name=name)
conv_1x1 = fluid.layers.conv2d(
conv1,
filter_size=1,
num_filters=out_dim // 2,
stride=stride,
param_attr=ParamAttr(
name=name + "_conv_1x1_weight", initializer=kaiming_init(conv1, 1)),
bias_attr=False,
name=name + '_conv_1x1')
conv_3x3 = fluid.layers.conv2d(
conv1,
filter_size=3,
num_filters=out_dim // 2,
stride=stride,
padding=1,
groups=out_dim // sr,
param_attr=ParamAttr(
name=name + "_conv_3x3_weight", initializer=kaiming_init(conv1, 3)),
bias_attr=False,
name=name + '_conv_3x3',
use_cudnn=False)
conv2 = fluid.layers.concat(
[conv_1x1, conv_3x3], axis=1, name=name + '_conv2')
pattr = ParamAttr(name=name + '_bn2_weight')
battr = ParamAttr(name=name + '_bn2_bias')
bn2 = fluid.layers.batch_norm(
input=conv2,
name=name + '_bn2',
param_attr=pattr,
bias_attr=battr,
moving_mean_name=name + '_bn2_running_mean',
moving_variance_name=name + '_bn2_running_var')
if stride == 1 and x.shape[1] == out_dim:
return fluid.layers.elementwise_add(
x=bn2, y=x, act='relu', name=name + "_add_relu")
else:
return fluid.layers.relu(bn2, name="_relu")
def make_layer(x, in_dim, out_dim, modules, block, name=None):
layers = block(x, out_dim, name=name + '_0')
for i in range(1, modules):
layers = block(layers, out_dim, name=name + '_' + str(i))
return layers
def make_hg_layer(x, in_dim, out_dim, modules, block, name=None):
layers = block(x, out_dim, stride=2, name=name + '_0')
for i in range(1, modules):
layers = block(layers, out_dim, name=name + '_' + str(i))
return layers
def make_layer_revr(x, in_dim, out_dim, modules, block, name=None):
for i in range(modules - 1):
x = block(x, in_dim, name=name + '_' + str(i))
layers = block(x, out_dim, name=name + '_' + str(modules - 1))
return layers
def make_unpool_layer(x, dim, name=None):
pattr = ParamAttr(name=name + '_weight', initializer=kaiming_init(x, 4))
battr = ParamAttr(name=name + '_bias', initializer=kaiming_init(x, 4))
layer = fluid.layers.conv2d_transpose(
input=x,
num_filters=dim,
filter_size=4,
stride=2,
padding=1,
param_attr=pattr,
bias_attr=battr)
return layer
@register
class Hourglass(object):
"""
Hourglass Network, see https://arxiv.org/abs/1603.06937
Args:
stack (int): stack of hourglass, 2 by default
dims (list): dims of each level in hg_module
modules (list): num of modules in each level
"""
__shared__ = ['stack']
def __init__(self,
stack=2,
dims=[256, 256, 384, 384, 512],
modules=[2, 2, 2, 2, 4],
block_name='fire'):
super(Hourglass, self).__init__()
self.stack = stack
assert len(dims) == len(modules), \
"Expected len of dims equal to len of modules, Receiced len of "\
"dims: {}, len of modules: {}".format(len(dims), len(modules))
self.dims = dims
self.modules = modules
self.num_level = len(dims) - 1
block_dict = {'fire': fire_block}
self.block = block_dict[block_name]
def __call__(self, input, name='hg'):
inter = self.pre(input, name + '_pre')
cnvs = []
for ind in range(self.stack):
hg = self.hg_module(
inter,
self.num_level,
self.dims,
self.modules,
name=name + '_hgs_' + str(ind))
cnv = _conv_norm(
hg,
3,
256,
bn_act='relu',
pad=1,
name=name + '_cnvs_' + str(ind))
cnvs.append(cnv)
if ind < self.stack - 1:
inter = _conv_norm(
inter, 1, 256, name=name + '_inters__' +
str(ind)) + _conv_norm(
cnv, 1, 256, name=name + '_cnvs__' + str(ind))
inter = fluid.layers.relu(inter)
inter = residual_block(
inter, 256, name=name + '_inters_' + str(ind))
return cnvs
def pre(self, x, name=None):
conv = _conv_norm(
x, 7, 128, stride=2, pad=3, bn_act='relu', name=name + '_0')
res1 = residual_block(conv, 256, stride=2, name=name + '_1')
res2 = residual_block(res1, 256, stride=2, name=name + '_2')
return res2
def hg_module(self,
x,
n=4,
dims=[256, 256, 384, 384, 512],
modules=[2, 2, 2, 2, 4],
make_up_layer=make_layer,
make_hg_layer=make_hg_layer,
make_low_layer=make_layer,
make_hg_layer_revr=make_layer_revr,
make_unpool_layer=make_unpool_layer,
name=None):
curr_mod = modules[0]
next_mod = modules[1]
curr_dim = dims[0]
next_dim = dims[1]
up1 = make_up_layer(
x, curr_dim, curr_dim, curr_mod, self.block, name=name + '_up1')
max1 = x
low1 = make_hg_layer(
max1, curr_dim, next_dim, curr_mod, self.block, name=name + '_low1')
low2 = self.hg_module(
low1,
n - 1,
dims[1:],
modules[1:],
make_up_layer=make_up_layer,
make_hg_layer=make_hg_layer,
make_low_layer=make_low_layer,
make_hg_layer_revr=make_hg_layer_revr,
make_unpool_layer=make_unpool_layer,
name=name + '_low2') if n > 1 else make_low_layer(
low1,
next_dim,
next_dim,
next_mod,
self.block,
name=name + '_low2')
low3 = make_hg_layer_revr(
low2, next_dim, curr_dim, curr_mod, self.block, name=name + '_low3')
up2 = make_unpool_layer(low3, curr_dim, name=name + '_up2')
merg = fluid.layers.elementwise_add(x=up1, y=up2, name=name + '_merg')
return merg