def Block(dim,
num_heads,
mlp_ratio=4.,
qkv_bias=False,
qk_scale=None,
drop=0.,
attn_drop=0.,
drop_path=0):
mlp_hidden_dim = int(dim * mlp_ratio)
return Sequential(
ShortCut(Identity(),
Sequential(
LayerNorm(),
Attention(dim,
num_heads=num_heads,
qkv_bias=qkv_bias,
qk_scale=qk_scale,
attn_drop=attn_drop,
proj_drop=drop), DropPath()),
mode='add'),
ShortCut(Identity(),
Sequential(
LayerNorm(),
Mlp(hidden_features=mlp_hidden_dim,
out_features=dim,
drop=drop), DropPath()),
mode='add'))
def __init__(self, *args, axis=1, activation='relu'):
"""
Parameters
----------
layer_defs : object
"""
super(ConcateBlock, self).__init__()
self.activation = get_activation(activation)
self.axis = axis
self.has_identity = False
for i in range(len(args)):
arg = args[i]
if isinstance(arg, (Layer, list, dict)):
if isinstance(arg, list):
arg = Sequential(*arg)
elif isinstance(arg, dict) and len(args) == 1:
for k, v in arg.items():
if isinstance(v, Identity):
self.has_identity = True
self.add_module('Identity', v)
else:
self.add_module(k, v)
elif isinstance(arg, dict) and len(args) > 1:
raise ValueError(
'more than one dict argument is not support.')
elif isinstance(arg, Identity):
self.has_identity = True
self.add_module('Identity', arg)
else:
self.add_module('branch{0}'.format(i + 1), arg)
if len(self._modules) == 1 and self.has_identity == False:
self.add_module('Identity', Identity())
self.to(self.device)
def basic_block(num_filters=64,base_width=64,strides=1,expansion = 4,conv_shortcut=False,use_bias=False,name=''):
shortcut = Identity()
if strides>1 or conv_shortcut is True:
shortcut =Conv2d_Block((1,1),num_filters=num_filters,strides=strides,auto_pad=True,padding_mode='zero',normalization='batch',activation=None,use_bias=use_bias)
return ShortCut2d(Sequential(Conv2d_Block((3,3),num_filters=num_filters,strides=strides,auto_pad=True,padding_mode='zero',normalization='batch',activation='relu',use_bias=use_bias),
Conv2d_Block((3,3),num_filters=num_filters,strides=1,auto_pad=True,padding_mode='zero',normalization='batch',activation=None,use_bias=use_bias)),
shortcut,activation='relu')
def se_bottleneck(num_filters=64,
strides=1,
expansion=4,
conv_shortcut=True,
use_bias=False,
name=''):
#width = int(num_filters * (base_width / 64.)) * 1#groups'
shortcut = Identity()
shortcut_name = 'Identity'
if strides > 1 or conv_shortcut is True:
shortcut = Conv2d_Block((1, 1),
num_filters=num_filters * expansion,
strides=strides,
auto_pad=True,
padding_mode='zero',
normalization='batch',
activation=None,
use_bias=use_bias,
name='downsample')
shortcut_name = 'downsample'
return ShortCut2d(
{
'branch1':
Sequential(
Conv2d_Block((1, 1),
num_filters=num_filters,
strides=strides,
auto_pad=True,
padding_mode='zero',
normalization='batch',
activation='relu',
use_bias=use_bias,
name='conv1'),
Conv2d_Block((3, 3),
num_filters=num_filters,
strides=1,
auto_pad=True,
padding_mode='zero',
normalization='batch',
activation='relu',
use_bias=use_bias,
name='conv2'),
Conv2d_Block((1, 1),
num_filters=num_filters * expansion,
strides=1,
auto_pad=True,
padding_mode='zero',
normalization='batch',
activation=None,
use_bias=use_bias,
name='conv3'),
SqueezeExcite(se_filters=num_filters // expansion,
num_filters=num_filters * expansion,
use_bias=True)),
shortcut_name:
shortcut
},
activation='relu')
def separable_resnet_block(num_filters=64, strides=1, activation='leaky_relu', normalization='batch',
use_spectral=False, dilation=1, name=''):
kernal = 1 if strides == 1 else 3
return [ShortCut2d(Sequential(
SeparableConv2d_Block((3, 3), depth_multiplier=1, strides=1, auto_pad=True, padding_mode='replicate',
use_spectral=use_spectral, normalization=normalization, activation=activation,
use_bias=False, dilation=dilation, name=name + '_0_conv'),
Conv2d_Block((1, 1), depth_multiplier=2, strides=1, auto_pad=True, padding_mode='replicate',
use_spectral=use_spectral, normalization=normalization, activation=activation, use_bias=False,
name=name + '_1_conv')), Identity(), activation=activation, name=name),
Conv2d_Block((kernal, kernal), num_filters=num_filters, strides=strides, auto_pad=True,
padding_mode='replicate', use_spectral=use_spectral, normalization=normalization,
activation=activation, use_bias=False, name=name + '_conv')]
def resblock_body(num_filters,
num_blocks,
all_narrow=True,
keep_output=False,
name=''):
return Sequential(
DarknetConv2D_BN_Mish((3, 3),
num_filters,
strides=(2, 2),
auto_pad=False,
padding=((1, 0), (1, 0)),
name=name + '_preconv1'),
ShortCut2d(
{
1:
DarknetConv2D_BN_Mish(
(1, 1),
num_filters // 2 if all_narrow else num_filters,
name=name + '_shortconv'),
0:
Sequential(
DarknetConv2D_BN_Mish(
(1, 1),
num_filters // 2 if all_narrow else num_filters,
name=name + '_mainconv'),
For(
range(num_blocks), lambda i: ShortCut2d(
Identity(),
Sequential(
DarknetConv2D_BN_Mish(
(1, 1),
num_filters // 2,
name=name + '_for{0}_1'.format(i)),
DarknetConv2D_BN_Mish(
(3, 3),
num_filters // 2
if all_narrow else num_filters,
name=name + '_for{0}_2'.format(i))),
mode='add')),
DarknetConv2D_BN_Mish(
(1, 1),
num_filters // 2 if all_narrow else num_filters,
name=name + '_postconv'))
},
mode='concate',
name=name + '_route'),
DarknetConv2D_BN_Mish((1, 1), num_filters, name=name + '_convblock5'))
def inverted_residual(in_filters,
num_filters=64,
strides=1,
expansion=4,
name=''):
mid_filters = int(round(in_filters * expansion))
layers = []
if expansion != 1:
layers.append(
Conv2d_Block((1, 1),
num_filters=mid_filters,
strides=1,
auto_pad=True,
padding_mode='zero',
normalization='batch',
activation='relu6'))
layers.append(
DepthwiseConv2d_Block((3, 3),
depth_multiplier=1,
strides=strides,
auto_pad=True,
padding_mode='zero',
normalization='batch',
activation='relu6'))
layers.append(
Conv2d_Block((1, 1),
num_filters=num_filters,
strides=1,
auto_pad=False,
padding_mode='zero',
normalization='batch',
activation=None))
if strides == 1 and in_filters == num_filters:
return ShortCut2d(Sequential(*layers), Identity(), activation=None)
else:
return Sequential(*layers)
def yolo4_body(anchors=None, num_classes=80, image_size=608):
anchors1 = generate_anchors(76) if anchors is None else anchors[0]
anchors2 = generate_anchors(38) if anchors is None else anchors[1]
anchors3 = generate_anchors(19) if anchors is None else anchors[2]
num_anchors = len(anchors1)
"""Create YOLO_V4 model CNN body in Pytorch."""
return Sequential(
DarknetConv2D_BN_Mish((3, 3), 32, name='first_layer'),
resblock_body(64, 1, all_narrow=False, name='block64'),
resblock_body(128, 2, name='block128'),
resblock_body(256, 8, name='block256'),
ShortCut2d(
{
1:
Sequential(
resblock_body(512, 8, name='block512'),
ShortCut2d(
{
1:
Sequential(
resblock_body(1024, 4, name='block1024'),
DarknetConv2D_BN_Leaky(
(1, 1), 512, name='pre_maxpool1'),
DarknetConv2D_BN_Leaky(
(3, 3), 1024, name='pre_maxpool2'),
DarknetConv2D_BN_Leaky(
(1, 1), 512, name='pre_maxpool3'),
ShortCut2d(
MaxPool2d((13, 13),
strides=(1, 1),
auto_pad=True),
MaxPool2d(
(9, 9), strides=(1, 1), auto_pad=True),
MaxPool2d(
(5, 5), strides=(1, 1), auto_pad=True),
Identity(),
mode='concate'),
DarknetConv2D_BN_Leaky(
(1, 1), 512, name='pre_y19_1'),
DarknetConv2D_BN_Leaky(
(3, 3), 1024, name='pre_y19_2'),
DarknetConv2D_BN_Leaky(
(1, 1), 512, name='y_19',
keep_output=True),
DarknetConv2D_BN_Leaky(
(1, 1), 256, name='pre_y19_upsample'),
Upsampling2d(scale_factor=2,
name='y19_upsample'),
),
0:
DarknetConv2D_BN_Leaky((1, 1), 256)
},
mode='concate'),
DarknetConv2D_BN_Leaky((1, 1), 256, name='pre_y38_1'),
DarknetConv2D_BN_Leaky((3, 3), 512, name='pre_y38_2'),
DarknetConv2D_BN_Leaky((1, 1), 256, name='pre_y38_3'),
DarknetConv2D_BN_Leaky((3, 3), 512, name='pre_y38_4'),
DarknetConv2D_BN_Leaky(
(1, 1), 256, name='y_38', keep_output=True),
DarknetConv2D_BN_Leaky(
(1, 1), 128, name='pre_y_38_upsample'),
Upsampling2d(scale_factor=2, name='y_38_upsample'),
),
0:
DarknetConv2D_BN_Leaky((1, 1), 128)
},
mode='concate'),
DarknetConv2D_BN_Leaky((1, 1), 128, name='pre_y76_concate1'),
DarknetConv2D_BN_Leaky((3, 3), 256, name='pre_y76_concate2'),
DarknetConv2D_BN_Leaky((1, 1), 128, name='pre_y76_concate3'),
DarknetConv2D_BN_Leaky((3, 3), 256, name='pre_y76_concate4'),
DarknetConv2D_BN_Leaky((1, 1), 128, name='pre_y76_concate5'),
ShortCut2d(
#y76_output
Sequential(DarknetConv2D_BN_Leaky((3, 3),
256,
name='pre_y76_output'),
DarknetConv2D((1, 1),
num_anchors * (num_classes + 5),
use_bias=True,
name='y76_output'),
YoloLayer(anchors=anchors1,
num_classes=num_classes,
grid_size=76,
img_dim=image_size),
name='y76_output'),
# y38_output
Sequential(
ShortCut2d(DarknetConv2D_BN_Leaky((3, 3),
256,
strides=(2, 2),
auto_pad=False,
padding=((1, 0), (1, 0)),
name='y76_downsample'),
branch_from='y_38',
mode='concate'),
DarknetConv2D_BN_Leaky((1, 1), 256, name='pre_y38_concate1'),
DarknetConv2D_BN_Leaky((3, 3), 512, name='pre_y38_concate2'),
DarknetConv2D_BN_Leaky((1, 1), 256, name='pre_y38_concate3'),
DarknetConv2D_BN_Leaky((3, 3), 512, name='pre_y38_concate4'),
DarknetConv2D_BN_Leaky((1, 1), 256, name='pre_y38_concate5'),
ShortCut2d(Sequential(DarknetConv2D_BN_Leaky(
(3, 3), 512, name='pre_y38_output'),
DarknetConv2D(
(1, 1),
num_anchors * (num_classes + 5),
use_bias=True,
name='y38_output'),
YoloLayer(anchors=anchors2,
num_classes=num_classes,
grid_size=38,
img_dim=image_size),
name='y38_output'),
Sequential(
ShortCut2d(DarknetConv2D_BN_Leaky(
(3, 3),
512,
strides=(2, 2),
auto_pad=False,
padding=((1, 0), (1, 0)),
name='y38_downsample'),
branch_from='y_19',
mode='concate'),
DarknetConv2D_BN_Leaky((1, 1),
512,
name='pre_y19_concate1'),
DarknetConv2D_BN_Leaky((3, 3),
1024,
name='pre_y19_concate2'),
DarknetConv2D_BN_Leaky((1, 1),
512,
name='pre_y19_concate3'),
DarknetConv2D_BN_Leaky((3, 3),
1024,
name='pre_y19_concate4'),
DarknetConv2D_BN_Leaky((1, 1),
512,
name='pre_y19_concate5'),
Sequential(DarknetConv2D_BN_Leaky(
(3, 3), 1024, name='pre_y19_output'),
DarknetConv2D(
(1, 1),
num_anchors * (num_classes + 5),
use_bias=True,
name='y19_output'),
YoloLayer(anchors=anchors3,
num_classes=num_classes,
grid_size=19,
img_dim=image_size),
name='y19_output')),
mode='concate')),
mode='concate'))
def efficient_block(expand_ratio=1,
filters_in=32,
filters_out=16,
kernel_size=3,
strides=1,
zero_pad=0,
se_ratio=0,
drop_rate=0.2,
is_shortcut=True,
name='',
**kwargs):
expand_ratio = kwargs.get('expand_ratio', expand_ratio)
is_shortcut = kwargs.get('id_skip', is_shortcut)
filters_in = kwargs.get('filters_in', filters_in)
filters_out = kwargs.get('filters_out', filters_out)
kernel_size = kwargs.get('kernel_size', kernel_size)
is_shortcut = filters_in == filters_out and strides == 1 and kwargs.get(
'id_skip', is_shortcut)
filters = filters_in * expand_ratio
if expand_ratio == 1 and strides == 1:
bottleneck = Sequential(
DepthwiseConv2d_Block((kernel_size, kernel_size),
depth_multiplier=1,
strides=strides,
auto_pad=True,
padding_mode='zero',
normalization='batch',
activation='swish'),
SqueezeExcite(se_filters=max(1, int(filters_in * se_ratio)),
num_filters=filters_in,
use_bias=True) if 0 < se_ratio <= 1 else Identity(),
Conv2d_Block((1, 1),
num_filters=filters_out,
strides=1,
auto_pad=True,
normalization='batch',
activation=None),
Dropout(dropout_rate=drop_rate)
if is_shortcut and drop_rate > 0 else Identity())
if is_shortcut:
return ShortCut2d(Identity(), bottleneck)
else:
return bottleneck
else:
bottleneck = Sequential(
Conv2d_Block((1, 1),
num_filters=filters,
strides=1,
auto_pad=True,
normalization='batch',
activation='swish'),
DepthwiseConv2d_Block((kernel_size, kernel_size),
depth_multiplier=1,
strides=strides,
auto_pad=True,
padding_mode='zero',
normalization='batch',
activation='swish'),
SqueezeExcite(se_filters=builtins.max(1, int(filters_in *
se_ratio)),
num_filters=filters,
use_bias=True) if 0 < se_ratio <= 1 else Identity(),
Conv2d_Block((1, 1),
num_filters=filters_out,
strides=1,
auto_pad=True,
normalization='batch',
activation=None),
Dropout(dropout_rate=drop_rate)
if is_shortcut and drop_rate > 0 else Identity())
if is_shortcut:
return ShortCut2d(Identity(), bottleneck)
else:
return bottleneck
def yolo4_body(num_classes=80, image_size=608):
anchors1 = to_tensor(np.array([12, 16, 19, 36, 40, 28]).reshape(-1, 2),
requires_grad=False)
anchors2 = to_tensor(np.array([36, 75, 76, 55, 72, 146]).reshape(-1, 2),
requires_grad=False)
anchors3 = to_tensor(np.array([142, 110, 192, 243, 459,
401]).reshape(-1, 2),
requires_grad=False)
num_anchors = len(anchors1)
"""Create YOLO_V4 model CNN body in Keras."""
return Sequential(
DarknetConv2D_BN_Mish((3, 3), 32, name='first_layer'),
resblock_body(64, 1, all_narrow=False, name='block64'),
resblock_body(128, 2, name='block128'),
resblock_body(256, 8, name='block256'),
ShortCut2d(
{
1:
Sequential(
resblock_body(512, 8, name='block512'),
ShortCut2d(
{
1:
Sequential(
resblock_body(1024, 4, name='block1024'),
DarknetConv2D_BN_Leaky(
(1, 1), 512, name='pre_maxpool1'),
DarknetConv2D_BN_Leaky(
(3, 3), 1024, name='pre_maxpool2'),
DarknetConv2D_BN_Leaky(
(1, 1), 512, name='pre_maxpool3'),
ShortCut2d(
MaxPool2d((13, 13),
strides=(1, 1),
auto_pad=True),
MaxPool2d(
(9, 9), strides=(1, 1), auto_pad=True),
MaxPool2d(
(5, 5), strides=(1, 1), auto_pad=True),
Identity(),
mode='concate'),
DarknetConv2D_BN_Leaky(
(1, 1), 512, name='pre_y19_1'),
DarknetConv2D_BN_Leaky(
(3, 3), 1024, name='pre_y19_2'),
DarknetConv2D_BN_Leaky(
(1, 1), 512, name='y_19',
keep_output=True),
DarknetConv2D_BN_Leaky(
(1, 1), 256, name='pre_y19_upsample'),
Upsampling2d(scale_factor=2,
name='y19_upsample'),
),
0:
DarknetConv2D_BN_Leaky((1, 1), 256)
},
mode='concate'),
DarknetConv2D_BN_Leaky((1, 1), 256, name='pre_y38_1'),
DarknetConv2D_BN_Leaky((3, 3), 512, name='pre_y38_2'),
DarknetConv2D_BN_Leaky((1, 1), 256, name='pre_y38_3'),
DarknetConv2D_BN_Leaky((3, 3), 512, name='pre_y38_4'),
DarknetConv2D_BN_Leaky(
(1, 1), 256, name='y_38', keep_output=True),
DarknetConv2D_BN_Leaky(
(1, 1), 128, name='pre_y_38_upsample'),
Upsampling2d(scale_factor=2, name='y_38_upsample'),
),
0:
DarknetConv2D_BN_Leaky((1, 1), 128)
},
mode='concate'),
DarknetConv2D_BN_Leaky((1, 1), 128, name='pre_y76_concate1'),
DarknetConv2D_BN_Leaky((3, 3), 256, name='pre_y76_concate2'),
DarknetConv2D_BN_Leaky((1, 1), 128, name='pre_y76_concate3'),
DarknetConv2D_BN_Leaky((3, 3), 256, name='pre_y76_concate4'),
DarknetConv2D_BN_Leaky((1, 1), 128, name='pre_y76_concate5'),
ShortCut2d(
#y76_output
Sequential(DarknetConv2D_BN_Leaky((3, 3),
256,
name='pre_y76_output'),
DarknetConv2D((1, 1),
num_anchors * (num_classes + 5),
use_bias=True,
name='y76_output'),
YoloLayer(anchors=anchors1,
num_classes=num_classes,
grid_size=76,
img_dim=image_size),
name='y76_output'),
# y38_output
Sequential(
ShortCut2d(DarknetConv2D_BN_Leaky((3, 3),
256,
strides=(2, 2),
auto_pad=False,
padding=((1, 0), (1, 0)),
name='y76_downsample'),
branch_from='y_38',
mode='concate'),
DarknetConv2D_BN_Leaky((1, 1), 256, name='pre_y38_concate1'),
DarknetConv2D_BN_Leaky((3, 3), 512, name='pre_y38_concate2'),
DarknetConv2D_BN_Leaky((1, 1), 256, name='pre_y38_concate3'),
DarknetConv2D_BN_Leaky((3, 3), 512, name='pre_y38_concate4'),
DarknetConv2D_BN_Leaky((1, 1), 256, name='pre_y38_concate5'),
ShortCut2d(Sequential(DarknetConv2D_BN_Leaky(
(3, 3), 512, name='pre_y38_output'),
DarknetConv2D(
(1, 1),
num_anchors * (num_classes + 5),
use_bias=True,
name='y38_output'),
YoloLayer(anchors=anchors2,
num_classes=num_classes,
grid_size=38,
img_dim=image_size),
name='y38_output'),
Sequential(
ShortCut2d(DarknetConv2D_BN_Leaky(
(3, 3),
512,
strides=(2, 2),
auto_pad=False,
padding=((1, 0), (1, 0)),
name='y38_downsample'),
branch_from='y_19',
mode='concate'),
DarknetConv2D_BN_Leaky((1, 1),
512,
name='pre_y19_concate1'),
DarknetConv2D_BN_Leaky((3, 3),
1024,
name='pre_y19_concate2'),
DarknetConv2D_BN_Leaky((1, 1),
512,
name='pre_y19_concate3'),
DarknetConv2D_BN_Leaky((3, 3),
1024,
name='pre_y19_concate4'),
DarknetConv2D_BN_Leaky((1, 1),
512,
name='pre_y19_concate5'),
Sequential(DarknetConv2D_BN_Leaky(
(3, 3), 1024, name='pre_y19_output'),
DarknetConv2D(
(1, 1),
num_anchors * (num_classes + 5),
use_bias=True,
name='y19_output'),
YoloLayer(anchors=anchors3,
num_classes=num_classes,
grid_size=19,
img_dim=image_size),
name='y19_output')),
mode='concate')),
mode='concate'))
def __init__(self,
*args,
axis=1,
branch_from=None,
activation=None,
mode='add',
name=None,
keep_output=False,
**kwargs):
"""
Args:
*args ():
axis ():
branch_from ():
activation ():
mode (str): 'add' 'dot' 'concate'
name (str):
keep_output (bool):
**kwargs ():
"""
super(ShortCut2d, self).__init__(name=name)
self.activation = get_activation(activation)
self.has_identity = False
self.mode = mode
self.axis = axis
self.branch_from = branch_from
self.branch_from_uuid = None
self.keep_output = keep_output
for i in range(len(args)):
arg = args[i]
if isinstance(arg, (Layer, torch.Tensor, list, dict)):
if isinstance(arg, list):
arg = Sequential(*arg)
elif isinstance(arg, OrderedDict) and len(args) == 1:
for k, v in arg.items():
if isinstance(v, Identity):
self.has_identity = True
self.add_module('Identity', v)
else:
self.add_module(k, v)
elif isinstance(arg, dict) and len(args) == 1:
keys = sorted(list(arg.keys()))
for k in keys:
v = arg[k]
if isinstance(v, Identity):
self.has_identity = True
self.add_module('Identity', v)
else:
self.add_module(str(k), v)
elif isinstance(arg, (dict, OrderedDict)) and len(args) > 1:
raise ValueError(
'more than one dict argument is not support.')
elif isinstance(arg, Identity):
self.has_identity = True
self.add_module('Identity', arg)
elif isinstance(arg, nn.Module):
if len(arg.name) > 0 and arg.name != arg.default_name:
self.add_module(arg.name, arg)
else:
self.add_module('branch{0}'.format(i + 1), arg)
else:
raise ValueError('{0} is not support.'.format(
arg.__class__.__name))
if len(
self._modules
) == 1 and self.has_identity == False and self.branch_from is None and mode != 'concate':
self.has_identity = True
self.add_module('Identity', Identity())
self.to(self.device)
def build_autoencoder(gan_role='discriminator',
use_skip_connections=False):
layers = OrderedDict()
ae = Sequential()
normalization = generator_norm if gan_role == 'generator' else discriminator_norm
build_block = generator_build_block if gan_role == 'generator' else discriminator_build_block
size_change = 4
initial_size = image_width // pow(2, size_change)
filter = 32
embedding_filters = 128
layers['head'] = Conv2d((3, 3),
32,
strides=1,
auto_pad=True,
use_bias=False,
activation=activation,
name='first_layer')
dilation = 1
# WIDTH 128=>64=>32=>16=>8=>
# FILTERS 32=>64=>64=>96=>96
for i in range(size_change):
if i % 2 == 1:
filter = filter + 32
if use_dilation:
dilation = 2 if i < 2 else 1
if build_block == BuildBlockMode.base.value:
layers['downsample_block{0}'.format(i)] = Conv2d_Block(
(3, 3),
filter,
strides=2,
auto_pad=True,
use_spectral=use_spectral,
use_bias=False,
activation=activation,
normalization=normalization,
name='base_block{0}'.format(i))
elif build_block == BuildBlockMode.resnet.value:
layers['downsample_block{0}'.format(i)] = Sequential(
resnet_block(filter,
strides=2,
activation=activation,
use_spectral=use_spectral,
normalization=normalization,
name='resnet_block{0}'.format(i)))
elif build_block == BuildBlockMode.bottleneck.value:
layers['downsample_block{0}'.format(i)] = bottleneck_block(
num_filters=filter,
strides=2,
reduce=4,
activation=activation,
use_spectral=use_spectral,
normalization=normalization,
name='bottleneck_block{0}'.format(i))
if use_self_attention and i == 2:
layers['self_attention'] = SelfAttention(16,
name='self_attention')
layers['embeddings'] = Conv2d((1, 1),
embedding_filters,
strides=1,
auto_pad=True,
use_bias=False,
activation=activation,
keep_output=True,
name='embeddings')
# 16=>32=>64=>128
# WIDTH 128<=64<=32<=16<=8
# FILTERS 32<=64<=64<=96<=96
for i in range(size_change):
if i % 2 == 0:
filter = filter - 32
if upsample_mode == UpsampleMode.transpose.value:
layers['upsampling{0}'.format(
size_change - i - 1)] = TransConv2d_Block(
(3, 3),
num_filters=filter,
strides=2,
auto_pad=True,
use_spectral=use_spectral,
use_bias=False,
activation=activation,
normalization=normalization,
dilation=dilation,
name='transconv_block{0}'.format(size_change - i - 1))
elif upsample_mode == UpsampleMode.pixel_shuffle.value:
layers['upsampling{0}'.format(
size_change - i - 1)] = Sequential(
Conv2d_Block((3, 3),
num_filters=4 * filter,
strides=1,
auto_pad=True,
use_spectral=use_spectral,
use_bias=False,
activation=activation,
normalization=normalization),
Upsampling2d(scale_factor=2,
mode=upsample_mode,
name='{0}{1}'.format(
upsample_mode, size_change - i - 1)))
else:
layers['upsampling{0}'.format(
size_change - i - 1)] = Sequential(
Conv2d_Block((3, 3),
num_filters=filter,
strides=1,
auto_pad=True,
use_spectral=use_spectral,
use_bias=False,
activation=activation,
normalization=normalization),
Upsampling2d(scale_factor=2,
mode=upsample_mode,
name='{0}{1}'.format(
upsample_mode, size_change - i - 1)))
if build_block == BuildBlockMode.base.value:
layers['upsampling_block{0}'.format(
size_change - i - 1)] = Conv2d_Block(
(3, 3),
filter,
strides=1,
auto_pad=True,
use_spectral=use_spectral,
use_bias=False,
activation=activation,
normalization=normalization,
dilation=dilation,
name='base_block{0}'.format(size_change - i - 1))
elif build_block == BuildBlockMode.resnet.value:
layers['upsampling_block{0}'.format(
size_change - i - 1)] = Sequential(
resnet_block(
filter,
strides=1,
activation=activation,
use_spectral=use_spectral,
normalization=normalization,
dilation=dilation,
name='resnet_block{0}'.format(size_change - i -
1)))
elif build_block == BuildBlockMode.bottleneck.value:
layers['upsampling_block{0}'.format(
size_change - i - 1)] = bottleneck_block(
filter,
strides=1,
activation=activation,
use_spectral=use_spectral,
normalization=normalization,
dilation=dilation,
name='resnet_block{0}'.format(size_change - i - 1))
elif build_block == BuildBlockMode.separable_resnet.value:
layers['upsampling_block{0}'.format(size_change - i - 1)] = (
separable_resnet_block(
filter,
strides=2,
activation=activation,
use_spectral=use_spectral,
normalization=discriminator_norm,
name='resnet_block{0}'.format(size_change - i - 1)))
layers['tail'] = Conv2d((3, 3),
generated_output_shape[0],
strides=1,
auto_pad=True,
use_bias=False,
activation='tanh',
name='last_layer')
if use_skip_connections:
shortcut3 = ShortCut2d(Identity(),
Sequential({
'downsample_block3':
layers['downsample_block3'],
'embeddings':
layers['embeddings'],
'upsampling3':
layers['upsampling3'],
}),
mode='add')
shortcut2dict = OrderedDict()
shortcut2dict['downsample_block2'] = layers['downsample_block2']
if use_self_attention:
shortcut2dict['self_attention'] = layers['self_attention']
shortcut2dict['shortcut3'] = shortcut3
shortcut2dict['upsampling_block3'] = layers['upsampling_block3']
shortcut2dict['upsampling2'] = layers['upsampling2']
shortcut2 = ShortCut2d(Identity(),
Sequential(shortcut2dict),
mode='add')
shortcut1 = ShortCut2d(Identity(),
Sequential({
'downsample_block1':
layers['downsample_block1'],
'shortcut2':
shortcut2,
'upsampling_block2':
layers['upsampling_block2'],
'upsampling1':
layers['upsampling1']
}),
mode='add')
shortcut0 = ShortCut2d(Identity(),
Sequential({
'downsample_block0':
layers['downsample_block0'],
'shortcut1':
shortcut1,
'upsampling_block1':
layers['upsampling_block1'],
'upsampling0':
layers['upsampling0']
}),
mode='add')
ae = Sequential({
'head': layers['head'],
'shortcut0': shortcut0,
'upsampling_block0': layers['upsampling_block0'],
'tail': layers['tail']
})
else:
ae = Sequential(layers)
return ae