def __init__(self,
in_channels,
out_channels,
kernel_size,
stride=1,
padding=0,
dilation=1,
groups=1,
bias='auto',
activation='relu',
inplace=True,
activate_last=True):
"""Init Conv Module with Normalization."""
super(ConvModule, self).__init__()
self.activation = activation
self.inplace = inplace
self.activate_last = activate_last
self.with_norm = True
self.with_activatation = activation is not None
if bias == 'auto':
bias = False if self.with_norm else True
self.with_bias = bias
self.conv = ops.Conv2d(in_channels,
out_channels,
kernel_size,
stride=stride,
padding=padding,
dilation=dilation,
groups=groups,
bias=bias)
self.in_channels = self.conv.in_channels
self.out_channels = self.conv.out_channels
self.kernel_size = self.conv.kernel_size
self.stride = self.conv.stride
self.padding = self.conv.padding
self.dilation = self.conv.dilation
self.transposed = self.conv.transposed
self.output_padding = self.conv.output_padding
self.groups = self.conv.groups
if self.with_norm:
norm_channels = out_channels if self.activate_last else in_channels
self.norm = ops.BatchNorm2d(norm_channels)
if self.with_activatation:
if self.activation not in ['relu']:
raise ValueError('{} is currently not supported.'.format(
self.activation))
if self.activation == 'relu':
self.activate = ops.Relu(inplace=inplace)
def conv(in_channel, out_channel, kernel_size=3, padding=None, sep=False):
"""Make a convolution layer with dilation 1, groups 1.
:param in_channel: number of input channels
:param out_channel: number of output channels
:param kernel_size: kernel size
:param padding: padding, Setting None to be same
:return: convolution layer as set
"""
if padding is None:
padding = kernel_size // 2
return ops.Conv2d(in_channel,
out_channel,
kernel_size=kernel_size,
padding=padding)
def __init__(self, init_plane):
"""Create InitialBlock layer.
:param init_plane: input channel.
:type init_plane: int
"""
super(InitialBlock, self).__init__()
self.conv = ops.Conv2d(in_channels=3,
out_channels=init_plane,
kernel_size=7,
stride=2,
padding=3,
bias=False)
self.batch = ops.BatchNorm2d(num_features=init_plane)
self.relu = ops.Relu()
self.maxpool2d = ops.MaxPool2d(kernel_size=3, stride=2, padding=1)
def conv33(in_channel, out_channel, stride=1, groups=1, bias=False):
"""Conv 3*3."""
if groups != 0 and in_channel % groups != 0:
raise ValueError(
'In channel "{}" is not a multiple of groups: "{}"'.format(
in_channel, groups))
if out_channel % groups != 0:
raise ValueError(
'Out channel "{}" is not a multiple of groups: "{}"'.format(
out_channel, groups))
return ops.Conv2d(in_channels=in_channel,
out_channels=out_channel,
kernel_size=3,
stride=stride,
bias=bias)
def __init__(self, inchannel, outchannel, expansion, stride=1):
"""Create ShortCut layer.
:param inchannel: input channel.
:type inchannel: int
:param outchannel: output channel.
:type outchannel: int
:param expansion: expansion
:type expansion: int
:param stride: the number to jump, default 1
:type stride: int
"""
super(ShortCut, self).__init__()
if stride != 1 or inchannel != outchannel * expansion:
self.conv1 = ops.Conv2d(in_channels=inchannel, out_channels=outchannel * expansion, kernel_size=1,
stride=stride, bias=False)
self.batch = ops.BatchNorm2d(num_features=outchannel * expansion)
def __init__(self, block_type, conv_num, growth_rate, type_prob, conv_prob,
growth_prob, G0, scale, code, architecture):
"""Construct the ESRN class.
:param net_desc: config of the searched structure
:type net_desc: list
"""
super(ESRN, self).__init__()
logging.info("start init ESRN")
self.arch = architecture
self.D = len(self.arch)
r = scale
G0 = G0
kSize = 3
n_colors = 3
self.SFENet1 = ops.Conv2d(n_colors,
G0,
kSize,
padding=(kSize - 1) // 2,
stride=1)
self.ERDBLayer = ERDBLayer(architecture, G0, kSize)
if r == 2 or r == 3:
self.UPNet = Sequential(
ops.Conv2d(G0,
G0 * 3,
kSize,
padding=(kSize - 1) // 2,
stride=1), ops.PixelShuffle(r),
ops.Conv2d(int(G0 * 3 / 4),
n_colors,
kSize,
padding=(kSize - 1) // 2,
stride=1))
elif r == 4:
self.UPNet = Sequential(
ops.Conv2d(G0,
G0 * 4,
kSize,
padding=(kSize - 1) // 2,
stride=1), ops.PixelShuffle(2),
ops.Conv2d(G0,
G0 * 4,
kSize,
padding=(kSize - 1) // 2,
stride=1), ops.PixelShuffle(2),
ops.Conv2d(G0,
n_colors,
kSize,
padding=(kSize - 1) // 2,
stride=1))
else:
raise ValueError("scale must be 2 or 3 or 4.")
def __init__(self,
C_in,
C_out,
kernel_size,
stride,
padding,
Conv2d='Conv2d',
affine=True,
use_relu6=False,
norm_layer='BN',
has_bn=True,
has_relu=True,
**kwargs):
"""Construct ConvBnRelu class."""
super(ConvBnRelu, self).__init__()
if Conv2d == 'Conv2d':
self.conv2d = ops.Conv2d(C_in,
C_out,
kernel_size,
stride=stride,
padding=padding,
bias=False)
elif Conv2d == 'ConvWS2d':
self.conv2d = ops.ConvWS2d(C_in,
C_out,
kernel_size,
stride=stride,
padding=padding,
bias=False)
if has_bn:
if norm_layer == 'BN':
self.batch_norm2d = ops.BatchNorm2d(C_out, affine=affine)
elif norm_layer == 'GN':
num_groups = kwargs.pop('num_groups')
self.batch_norm2d = ops.GroupNorm(num_groups,
C_out,
affine=affine)
elif norm_layer == 'Sync':
self.batch_norm2d = ops.SyncBatchNorm(C_out, affine=affine)
if has_relu:
if use_relu6:
self.relu = ops.Relu6(inplace=False)
else:
self.relu = ops.Relu(inplace=False)
def __init__(self,
in_channels,
out_channels,
kernel_size,
stride,
padding,
dilation=1,
groups=1,
bias=False):
super(BN_Conv2d, self).__init__()
self.seq = Sequential(
ops.Conv2d(in_channels,
out_channels,
kernel_size=kernel_size,
stride=stride,
padding=padding,
dilation=dilation,
groups=groups,
bias=bias), ops.BatchNorm2d(out_channels), ops.Relu())
def __init__(self,
C_in,
C_out,
kernel_size,
stride,
padding,
affine=True,
use_relu6=False):
"""Construct ConvBnRelu class."""
super(ConvBnRelu, self).__init__()
self.conv2d = ops.Conv2d(C_in,
C_out,
kernel_size,
stride=stride,
padding=padding,
bias=False)
self.batch_norm2d = ops.BatchNorm2d(C_out, affine=affine)
if use_relu6:
self.relu = ops.Relu6(inplace=False)
else:
self.relu = ops.Relu(inplace=False)
def make_res_layer(block,
inplanes,
planes,
blocks,
stride=1,
dilation=1,
style='pytorch',
with_cp=False):
"""Build resnet layer."""
downsample = None
if stride != 1 or inplanes != planes * block.expansion:
conv_layer = ops.Conv2d(inplanes,
planes * block.expansion,
kernel_size=1,
stride=stride,
bias=False)
norm_layer = ops.BatchNorm2d(planes * block.expansion)
downsample = Sequential(conv_layer, norm_layer)
layers = []
layers.append(
block(inplanes=inplanes,
planes=planes,
stride=stride,
dilation=dilation,
downsample=downsample,
style=style,
with_cp=with_cp))
inplanes = planes * block.expansion
for i in range(1, blocks):
layers.append(
block(inplanes=inplanes,
planes=planes,
stride=1,
dilation=dilation,
style=style,
with_cp=with_cp))
return Sequential(*layers)
def conv1X1(inchannel, outchannel, stride=1):
"""Create conv1X1 layer."""
return ops.Conv2d(inchannel, outchannel, kernel_size=1, stride=stride, bias=False)
def __init__(self, C_in, C_out, kernel_size, stride, padding, affine=True):
"""Init ReLUConvBN."""
super(ReLUConvBN, self).__init__()
self.relu = ops.Relu(inplace=False)
self.conv = ops.Conv2d(C_in, C_out, kernel_size, stride=stride, padding=padding, bias=False)
self.bn = ops.BatchNorm2d(C_out, affine=affine)
'none': lambda C, stride, affine, repeats=1: ops.Zero(stride),
'avg_pool_3x3': lambda C, stride, affine, repeats=1: ops.AvgPool2d(
3, stride=stride, padding=1, count_include_pad=False),
'max_pool_3x3': lambda C, stride, affine, repeats=1: ops.MaxPool2d(
3, stride=stride, padding=1),
'global_average_pool': lambda C, stride, affine, repeats=1: Seq(GAPConv1x1(C, C)),
'skip_connect': lambda C, stride, affine, repeats=1: ops.Identity() if stride == 1 else FactorizedReduce(
C, C, affine=affine),
'sep_conv_3x3': lambda C, stride, affine, repeats=1: SeparatedConv(C, C, 3, stride, 1, affine=affine),
'sep_conv_5x5': lambda C, stride, affine, repeats=1: SeparatedConv(C, C, 5, stride, 2, affine=affine),
'sep_conv_7x7': lambda C, stride, affine, repeats=1: SeparatedConv(C, C, 7, stride, 3, affine=affine),
'dil_conv_3x3': lambda C, stride, affine, repeats=1: DilConv(C, C, 3, stride, 2, 2, affine=affine),
'dil_conv_5x5': lambda C, stride, affine, repeats=1: DilConv(C, C, 5, stride, 4, 2, affine=affine),
'conv_7x1_1x7': lambda C, stride, affine, repeats=1: Seq(
ops.Relu(inplace=False),
ops.Conv2d(C, C, (1, 7), stride=(1, stride), padding=(0, 3), bias=False),
ops.Conv2d(C, C, (7, 1), stride=(stride, 1), padding=(3, 0), bias=False),
ops.BatchNorm2d(C, affine=affine)),
'conv1x1': lambda C, stride, affine, repeats=1: Seq(
conv1X1(C, C, stride=stride),
ops.BatchNorm2d(C, affine=affine),
ops.Relu(inplace=False)),
'conv3x3': lambda C, stride, affine, repeats=1: Seq(
conv3x3(C, C, stride=stride),
ops.BatchNorm2d(C, affine=affine),
ops.Relu(inplace=False)),
'conv5x5': lambda C, stride, affine, repeats=1: Seq(
conv5x5(C, C, stride=stride),
ops.BatchNorm2d(C, affine=affine),
ops.Relu(inplace=False)),
'conv7x7': lambda C, stride, affine, repeats=1: Seq(
def conv5x5(inchannel, outchannel, stride=1, bias=False, dilation=1):
"""Create Convolution 5x5."""
return ops.Conv2d(inchannel, outchannel, kernel_size=5, stride=stride,
padding=2, dilation=dilation, bias=bias)
def _make_stem_layer(self):
"""Make stem layer."""
self.conv1 = ops.Conv2d(3, self.inplanes, kernel_size=7, stride=2, padding=3, bias=False)
self.norm1 = ops.BatchNorm2d(64)
self.relu = ops.Relu(inplace=True)
self.maxpool = ops.MaxPool2d(kernel_size=3, stride=2, padding=1)
def conv7x7(inchannel, outchannel, stride=1, bias=False, dilation=1):
"""Create Convolution 7x7."""
return ops.Conv2d(inchannel, outchannel, kernel_size=7, stride=stride,
padding=3, dilation=dilation, bias=bias)
def conv11(in_channel, out_channel, stride=1, bias=False):
"""Conv 1*1."""
return ops.Conv2d(in_channels=in_channel, out_channels=out_channel, kernel_size=1,
stride=stride, bias=bias)
def __init__(self, init_channels, stem_multi):
"""Init PreOneStem."""
super(PreOneStem, self).__init__()
self._c_curr = init_channels * stem_multi
self.conv2d = ops.Conv2d(3, self._c_curr, 3, padding=1, bias=False)
self.batchNorm2d = ops.BatchNorm2d(self._c_curr)
def conv3x3(inchannel, outchannel, groups=1, stride=1, bias=False, dilation=1):
"""Create conv3x3 layer."""
return ops.Conv2d(inchannel, outchannel, kernel_size=3, stride=stride,
padding=dilation, groups=groups, bias=bias, dilation=dilation)
