def __init__(self, ctx=mx.cpu(), warmup=5, runs=25, inputs=None):
# Set the default Inputs
default_parameters = {
"data": (32, 3, 256),
"data_initializer": nd.normal,
"pool_size": 2,
"strides": None,
"padding": 0,
"layout": "NCW",
"run_backward": True,
"dtype": "float32"
}
super().__init__(ctx=ctx,
warmup=warmup,
runs=runs,
default_parameters=default_parameters,
custom_parameters=inputs)
self.data = get_mx_ndarray(ctx=self.ctx,
in_tensor=self.inputs["data"],
dtype=self.inputs["dtype"],
initializer=self.inputs["data_initializer"],
attach_grad=self.inputs["run_backward"])
self.block = nn.AvgPool1D(pool_size=self.inputs["pool_size"],
strides=self.inputs["strides"],
padding=self.inputs["padding"],
layout=self.inputs["layout"])
self.block.initialize(ctx=self.ctx)
def test_pool():
layers1d = [
nn.MaxPool1D(),
nn.MaxPool1D(3),
nn.MaxPool1D(3, 2),
nn.AvgPool1D(),
nn.AvgPool1D(count_include_pad=False),
nn.GlobalAvgPool1D(),
]
for layer in layers1d:
check_layer_forward(layer, (1, 2, 10))
layers2d = [
nn.MaxPool2D(),
nn.MaxPool2D((3, 3)),
nn.MaxPool2D(3, 2),
nn.AvgPool2D(),
nn.AvgPool2D(count_include_pad=False),
nn.GlobalAvgPool2D(),
]
for layer in layers2d:
check_layer_forward(layer, (1, 2, 10, 10))
layers3d = [
nn.MaxPool3D(),
nn.MaxPool3D((3, 3, 3)),
nn.MaxPool3D(3, 2),
nn.AvgPool3D(),
nn.AvgPool3D(count_include_pad=False),
nn.GlobalAvgPool3D(),
]
for layer in layers3d:
check_layer_forward(layer, (1, 2, 10, 10, 10))
# test ceil_mode
x = mx.nd.zeros((2, 2, 10, 10))
layer = nn.MaxPool2D(3, ceil_mode=False)
layer.collect_params().initialize()
assert (layer(x).shape==(2, 2, 3, 3))
layer = nn.MaxPool2D(3, ceil_mode=True)
layer.collect_params().initialize()
assert (layer(x).shape==(2, 2, 4, 4))
def _make_layer(self,
stage_index,
block,
planes,
blocks,
strides=1,
dilation=1,
pre_dilation=1,
avg_down=False,
norm_layer=None,
last_gamma=False,
dropblock_prob=0,
input_size=224,
use_splat=False,
avd=False):
downsample = None
if strides != 1 or self.inplanes != planes * block.expansion:
downsample = nn.HybridSequential(prefix='down%d_' % stage_index)
with downsample.name_scope():
if avg_down:
if pre_dilation == 1:
downsample.add(
nn.AvgPool1D(pool_size=strides,
strides=strides,
ceil_mode=True,
count_include_pad=False))
elif strides == 1:
downsample.add(
nn.AvgPool1D(pool_size=1,
strides=1,
ceil_mode=True,
count_include_pad=False))
else:
downsample.add(
nn.AvgPool1D(pool_size=pre_dilation * strides,
strides=strides,
padding=1,
ceil_mode=True,
count_include_pad=False))
downsample.add(
nn.Conv1D(channels=planes * block.expansion,
kernel_size=1,
strides=1,
use_bias=False,
in_channels=self.inplanes))
downsample.add(
norm_layer(in_channels=planes * block.expansion,
**self.norm_kwargs))
else:
downsample.add(
nn.Conv1D(channels=planes * block.expansion,
kernel_size=1,
strides=strides,
use_bias=False,
in_channels=self.inplanes))
downsample.add(
norm_layer(in_channels=planes * block.expansion,
**self.norm_kwargs))
layers = nn.HybridSequential(prefix='layers%d_' % stage_index)
with layers.name_scope():
if dilation in (1, 2):
layers.add(
block(planes,
cardinality=self.cardinality,
bottleneck_width=self.bottleneck_width,
strides=strides,
dilation=pre_dilation,
downsample=downsample,
previous_dilation=dilation,
norm_layer=norm_layer,
norm_kwargs=self.norm_kwargs,
last_gamma=last_gamma,
dropblock_prob=dropblock_prob,
input_size=input_size,
use_splat=use_splat,
avd=avd,
avd_first=self.avd_first,
radix=self.radix,
in_channels=self.inplanes,
split_drop_ratio=self.split_drop_ratio))
elif dilation == 4:
layers.add(
block(planes,
cardinality=self.cardinality,
bottleneck_width=self.bottleneck_width,
strides=strides,
dilation=pre_dilation,
downsample=downsample,
previous_dilation=dilation,
norm_layer=norm_layer,
norm_kwargs=self.norm_kwargs,
last_gamma=last_gamma,
dropblock_prob=dropblock_prob,
input_size=input_size,
use_splat=use_splat,
avd=avd,
avd_first=self.avd_first,
radix=self.radix,
in_channels=self.inplanes,
split_drop_ratio=self.split_drop_ratio))
else:
raise RuntimeError(
"=> unknown dilation size: {}".format(dilation))
input_size = _update_input_size(input_size, strides)
self.inplanes = planes * block.expansion
for i in range(1, blocks):
layers.add(
block(planes,
cardinality=self.cardinality,
bottleneck_width=self.bottleneck_width,
dilation=dilation,
previous_dilation=dilation,
norm_layer=norm_layer,
norm_kwargs=self.norm_kwargs,
last_gamma=last_gamma,
dropblock_prob=dropblock_prob,
input_size=input_size,
use_splat=use_splat,
avd=avd,
avd_first=self.avd_first,
radix=self.radix,
in_channels=self.inplanes,
split_drop_ratio=self.split_drop_ratio))
return layers
def __init__(self,
channels,
cardinality=1,
bottleneck_width=64,
strides=1,
dilation=1,
downsample=None,
previous_dilation=1,
norm_layer=None,
norm_kwargs=None,
last_gamma=False,
dropblock_prob=0,
input_size=None,
use_splat=False,
radix=2,
avd=False,
avd_first=False,
in_channels=None,
split_drop_ratio=0,
**kwargs):
super(Bottleneck, self).__init__()
group_width = int(channels * (bottleneck_width / 64.)) * cardinality
norm_kwargs = norm_kwargs if norm_kwargs is not None else {}
self.dropblock_prob = dropblock_prob
self.use_splat = use_splat
self.avd = avd and (strides > 1 or previous_dilation != dilation)
self.avd_first = avd_first
if self.dropblock_prob > 0:
self.dropblock1 = DropBlock(dropblock_prob, 3, group_width,
*input_size)
if self.avd:
if avd_first:
input_size = _update_input_size(input_size, strides)
self.dropblock2 = DropBlock(dropblock_prob, 3, group_width,
*input_size)
if not avd_first:
input_size = _update_input_size(input_size, strides)
else:
input_size = _update_input_size(input_size, strides)
self.dropblock2 = DropBlock(dropblock_prob, 3, group_width,
*input_size)
self.dropblock3 = DropBlock(dropblock_prob, 3, channels * 4,
*input_size)
self.conv1 = nn.Conv1D(channels=group_width,
kernel_size=1,
use_bias=False,
in_channels=in_channels)
self.bn1 = norm_layer(in_channels=group_width, **norm_kwargs)
self.relu1 = nn.Activation('relu')
if self.use_splat:
self.conv2 = SplitAttentionConv(channels=group_width,
kernel_size=3,
strides=1 if self.avd else strides,
padding=dilation,
dilation=dilation,
groups=cardinality,
use_bias=False,
in_channels=group_width,
norm_layer=norm_layer,
norm_kwargs=norm_kwargs,
radix=radix,
drop_ratio=split_drop_ratio,
**kwargs)
else:
self.conv2 = nn.Conv1D(channels=group_width,
kernel_size=3,
strides=1 if self.avd else strides,
padding=dilation,
dilation=dilation,
groups=cardinality,
use_bias=False,
in_channels=group_width,
**kwargs)
self.bn2 = norm_layer(in_channels=group_width, **norm_kwargs)
self.relu2 = nn.Activation('relu')
self.conv3 = nn.Conv1D(channels=channels * 4,
kernel_size=1,
use_bias=False,
in_channels=group_width)
if not last_gamma:
self.bn3 = norm_layer(in_channels=channels * 4, **norm_kwargs)
else:
self.bn3 = norm_layer(in_channels=channels * 4,
gamma_initializer='zeros',
**norm_kwargs)
if self.avd:
self.avd_layer = nn.AvgPool1D(3, strides, padding=1)
self.relu3 = nn.Activation('relu')
self.downsample = downsample
self.dilation = dilation
self.strides = strides
