def alpha_zero_resnet(n_labels=4992,
channels=256,
channels_value_head=1,
channels_policy_head=81,
num_res_blocks=19,
value_fc_size=256,
bn_mom=0.9,
act_type='relu'):
sym = alpha_zero_symbol(
num_filter=channels,
channels_value_head=channels_value_head,
channels_policy_head=channels_policy_head,
workspace=1024,
value_fc_size=value_fc_size,
num_res_blocks=num_res_blocks,
bn_mom=bn_mom,
n_labels=n_labels,
grad_scale_policy=1.,
grad_scale_value=1.,
)
"""
Constructs the alpha zero network for gluon representation using a symbol block
"""
net = SymbolBlock(outputs=sym, inputs=mx.sym.var('data'))
return net
def expand_network(network, layers, num_filters, min_num_filter=128):
output_layers, inputs, params = parse_network(network, layers)
layer = output_layers[-1]
for idx, num_filter in enumerate(num_filters):
# num_filter_1x1 = max(min_num_filter, num_filter//2)
num_filter_1x1 = max(min_num_filter, num_filter)
layer = conv_act_layer(data=layer,
num_filter=num_filter_1x1,
kernel=(1, 1),
stride=(1, 1),
pad=(0, 0),
name='multi_feat_{}_conv_1x1'.format(idx))
layer = conv_act_layer(data=layer,
num_filter=num_filter,
kernel=(3, 3),
stride=(2, 2),
pad=(1, 1),
name='multi_feat_{}_conv_3x3'.format(idx))
output_layers.append(layer)
output_blocks = SymbolBlock(outputs=output_layers,
inputs=inputs,
params=params)
return output_blocks
def func(pretrained=False,
tag=None,
root='~/.mxnet/models',
ctx=cpu(0),
**kwargs):
r"""Quantized model.
Parameters
----------
pretrained : bool or str
Boolean value controls whether to load the default pretrained weights for model.
String value represents the hashtag for a certain version of pretrained weights.
tag : str, default is None
Optional length-8 sha1sum of parameter file. If `None`, best parameter file
will be used.
ctx : Context, default CPU
The context in which to load the pretrained weights.
root : str, default $MXNET_HOME/models
Location for keeping the model parameters.
"""
from ..model_zoo import get_model
from ..model_store import get_model_file
curr_dir = os.path.abspath(os.path.dirname(__file__))
model_name = name.replace('mobilenet1_', 'mobilenet1.')
model_name = model_name.replace('mobilenet0_', 'mobilenet0.')
json_file = os.path.join(curr_dir, '{}-symbol.json'.format(model_name))
base_name = '_'.join(model_name.split('_')[:-1])
with warnings.catch_warnings(record=True) as w:
warnings.simplefilter("always")
param_file = get_model_file(base_name, tag=tag,
root=root) if pretrained else None
net = get_model('_'.join(model_name.split('_')[:-1]),
prefix=sym_prefix)
classes = getattr(net, 'classes', [])
sym_net = SymbolBlock.imports(json_file, ['data'], None, ctx=ctx)
if param_file:
# directly imports weights saved by save_parameters is not applicable
# so we hack it by load and export once to a temporary params file
import tempfile
net.load_params(param_file)
net.hybridize()
if '512' in base_name:
net(mx.nd.zeros((1, 3, 512, 512)))
elif '300' in base_name:
net(mx.nd.zeros((1, 3, 300, 300)))
else:
net(mx.nd.zeros((1, 3, 224, 224)))
with tempfile.TemporaryDirectory() as tmpdirname:
prefix = os.path.join(tmpdirname, 'tmp')
net.export(prefix, epoch=0)
param_prefix = prefix + '-0000.params'
sym_net.collect_params().load(param_prefix)
net.classes = classes
net.reset_class = _not_impl
net.set_nms = _not_impl
return net
def check_amp_convert_hybrid_block():
# Test conversion for hybrid block on CPU
model_cpu = get_model("resnet50_v1")
model_cpu.collect_params().initialize(ctx=mx.cpu())
model_cpu.hybridize()
model_cpu(mx.nd.random.uniform(0, 1, shape=(1, 3, 224, 224), ctx=mx.cpu()))
converted_model_cpu = amp.convert_hybrid_block(model_cpu)
# Test with real world model, default inputs for convert_hybrid_block
model = get_model("resnet50_v1")
model.collect_params().initialize(ctx=mx.gpu())
model.hybridize()
model(mx.nd.zeros((1, 3, 224, 224)))
converted_model = amp.convert_hybrid_block(model)
result = converted_model.forward(mx.nd.zeros((1, 3, 224, 224),
dtype=np.float32))
result = converted_model.forward(mx.nd.zeros((1, 3, 224, 224),
dtype=np.float32))
# Test with real world model, tweak inputs for convert_hybrid_block
converted_model = amp.convert_hybrid_block(model, target_dtype="float16",
target_dtype_ops=["Convolution"])
result = converted_model.forward(mx.nd.zeros((1, 3, 224, 224),
dtype=np.float32))
result = converted_model.forward(mx.nd.zeros((1, 3, 224, 224),
dtype=np.float32))
# Check symbolic block
dir_path = os.path.dirname(os.path.realpath(__file__))
model_path = os.path.join(dir_path, 'model')
if not os.path.isdir(model_path):
os.mkdir(model_path)
prefix, epoch = download_model("imagenet1k-resnet-18", dst_dir=model_path)
net = SymbolBlock.imports(os.path.join(model_path, "imagenet1k-resnet-18-symbol.json"),
input_names=["data", "softmax_label"],
param_file=os.path.join(model_path, "imagenet1k-resnet-18-0000.params"))
net.collect_params().reset_ctx(ctx=mx.gpu())
net.hybridize()
net(mx.nd.zeros((1, 3, 224, 224)), mx.nd.zeros((1,)))
converted_model = amp.convert_hybrid_block(net)
result = converted_model.forward(mx.nd.zeros((1, 3, 224, 224)), mx.nd.zeros((1,)))
result = converted_model.forward(mx.nd.zeros((1, 3, 224, 224)), mx.nd.zeros((1,)))
# Check symbolic block, tweaked inputs
converted_model = amp.convert_hybrid_block(net, target_dtype="float16", target_dtype_ops=["Convolution"])
result = converted_model.forward(mx.nd.zeros((1, 3, 224, 224)), mx.nd.zeros((1, )))
result = converted_model.forward(mx.nd.zeros((1, 3, 224, 224)), mx.nd.zeros((1, )))
params = converted_model.collect_params()
assert params["stage2_unit1_conv2_weight"].dtype == np.float32
# Pass cast_optional_params as True to convert_hybrid_block
converted_model = amp.convert_hybrid_block(net, target_dtype="float16", target_dtype_ops=["Convolution"],
cast_optional_params=True)
params = converted_model.collect_params()
assert params["stage2_unit1_conv2_weight"].dtype == np.float16
def func(pretrained=False, tag=None, root='~/.mxnet/models', ctx=cpu(0), **kwargs):
r"""Quantized model.
Parameters
----------
pretrained : bool or str
Boolean value controls whether to load the default pretrained weights for model.
String value represents the hashtag for a certain version of pretrained weights.
tag : str, default is None
Optional length-8 sha1sum of parameter file. If `None`, best parameter file
will be used.
ctx : Context, default CPU
The context in which to load the pretrained weights.
root : str, default $MXNET_HOME/models
Location for keeping the model parameters.
"""
from ..model_zoo import get_model
from ..model_store import get_model_file
curr_dir = os.path.abspath(os.path.dirname(__file__))
model_name = name.replace('mobilenet1_', 'mobilenet1.')
model_name = model_name.replace('mobilenet0_', 'mobilenet0.')
json_file = os.path.join(curr_dir, '{}-symbol.json'.format(model_name))
base_name = '_'.join(model_name.split('_')[:-1])
with warnings.catch_warnings(record=True) as w:
warnings.simplefilter("always")
param_file = get_model_file(base_name, tag=tag, root=root) if pretrained else None
net = get_model('_'.join(model_name.split('_')[:-1]), prefix=sym_prefix)
classes = getattr(net, 'classes', [])
sym_net = SymbolBlock.imports(json_file, ['data'], None, ctx=ctx)
if param_file:
# directly imports weights saved by save_parameters is not applicable
# so we hack it by load and export once to a temporary params file
import tempfile
net.load_params(param_file)
net.hybridize()
if '512' in base_name:
net(mx.nd.zeros((1, 3, 512, 512)))
elif '300' in base_name:
net(mx.nd.zeros((1, 3, 300, 300)))
else:
net(mx.nd.zeros((1, 3, 224, 224)))
with tempfile.TemporaryDirectory() as tmpdirname:
prefix = os.path.join(tmpdirname, 'tmp')
net.export(prefix, epoch=0)
param_prefix = prefix + '-0000.params'
sym_net.collect_params().load(param_prefix)
sym_net.classes = classes
sym_net.reset_class = _not_impl
sym_net.set_nms = _not_impl
return sym_net
def load_net(self):
net_cl = SymbolBlock.imports(self.net_path_cl, ['data'],
self.params_path_cl)
net_tl = SymbolBlock.imports(self.net_path_tl, ['data'],
self.params_path_tl)
return net_cl, net_tl
def network_extractor(network, layers):
output_layers, inputs, params = parse_network(network, layers)
output_blocks = SymbolBlock(outputs=output_layers,
inputs=inputs,
params=params)
return output_blocks
def quantize(self, tune_cfg, model, dataloader, q_func=None):
"""The function is used to do MXNet calibration and quanitization in post-training
quantization.
Args:
tune_cfg (dict): quantization config.
model (object): model need to do quantization.
dataloader (object): calibration dataset.
q_func (optional): training function for quantization aware training mode,
unimplement yet for MXNet.
Returns:
(dict): quantized model
"""
assert q_func is None, "quantization aware training mode is not support on mxnet"
# get symbol from FP32 model
if isinstance(model.model, mx.gluon.HybridBlock):
# transfer hybridblock to symbol
sym, arg_params, aux_params, calib_data = \
self._get_gluon_symbol(model.model, dataloader=dataloader)
data_names = [pair[0] for pair in calib_data.provide_data]
self.__config_dict['calib_data'] = calib_data
elif isinstance(model.model[0], mx.symbol.Symbol):
sym, arg_params, aux_params = model.model
self.__config_dict['calib_data'] = dataloader
else:
raise ValueError(
'Need a symbol model or HybridBlock model, while received %s' % str(
type(model.model)))
self._cfg_to_qconfig(tune_cfg)
self.th_dict = None
qconfig = self.__config_dict
sym = self._get_backedn_graph(sym, qconfig['ctx'])
# 1. quantize_symbol
if _check_version(mx.__version__, '1.7.0'):
qsym, calib_layer = mx.contrib.quantization._quantize_symbol(
sym, qconfig['ctx'],
excluded_symbols=qconfig['excluded_sym_names'],
excluded_operators=qconfig['excluded_op_names'],
offline_params=list(arg_params.keys()),
quantized_dtype=qconfig['quantized_dtype'],
quantize_mode=qconfig['quantize_mode'],
quantize_granularity=qconfig['quantize_granularity'])
else:
qsym, calib_layer = mx.contrib.quantization._quantize_symbol(
sym, qconfig['ctx'],
excluded_symbols=qconfig['excluded_sym_names'],
excluded_operators=qconfig['excluded_op_names'],
offline_params=list(arg_params.keys()),
quantized_dtype=qconfig['quantized_dtype'],
quantize_mode=qconfig['quantize_mode'],)
# 2. Do calibration to get th_dict
th_dict = self._get_calibration_th(
sym, arg_params, aux_params, calib_layer, qconfig)
self.th_dict = th_dict
# 3. set th_dict to quantized symbol
qsym = mx.contrib.quantization._calibrate_quantized_sym(qsym, th_dict)
# 4. quantize params
qarg_params = mx.contrib.quantization._quantize_params(
qsym, arg_params, th_dict)
qsym = self._get_backedn_graph(qsym, qconfig['ctx'])
if isinstance(model.model, mx.gluon.HybridBlock):
from mxnet.gluon import SymbolBlock
data_sym = []
for name in data_names:
data_sym.append(mx.sym.var(name))
net = SymbolBlock(qsym, data_sym)
with TemporaryDirectory() as tmpdirname:
prefix = os.path.join(tmpdirname, 'tmp')
param_name = '%s-%04d.params' % (prefix + 'net-quantized', 0)
save_dict = {('arg:%s' % k): v.as_in_context(mx.cpu())
for k, v in qarg_params.items()}
save_dict.update({('aux:%s' % k): v.as_in_context(mx.cpu())
for k, v in aux_params.items()})
mx.ndarray.save(param_name, save_dict) # pylint: disable=no-member
net.collect_params().load(param_name, cast_dtype=True, dtype_source='saved')
net.collect_params().reset_ctx(self.__config_dict['ctx'])
return Model(net)
return Model((qsym, qarg_params, aux_params))
def __init__(self, **kwargs):
super(InceptionBNStages, self).__init__(**kwargs)
# data
data = mx.symbol.Variable(name="data")
# stage 1
conv1 = get_conv(data=data,
num_filter=64,
kernel=(7, 7),
stride=(2, 2),
pad=(3, 3),
name='1')
pool1 = mx.symbol.Pooling(data=conv1,
kernel=(3, 3),
stride=(2, 2),
name='pool_1',
pool_type='max')
conv2red = get_conv(data=pool1,
num_filter=64,
kernel=(1, 1),
stride=(1, 1),
name='2_red')
conv2 = get_conv(data=conv2red,
num_filter=192,
kernel=(3, 3),
stride=(1, 1),
pad=(1, 1),
name='2')
# stage 2
pool2 = mx.symbol.Pooling(data=data,
kernel=(3, 3),
stride=(2, 2),
name='pool_2',
pool_type='max')
in3a = get_inception_a(pool2, 64, 64, 64, 64, 96, "avg", 32, '3a')
in3b = get_inception_a(in3a, 64, 64, 96, 64, 96, "avg", 64, '3b')
in3c = get_inception_b(in3b, 128, 160, 64, 96, '3c')
# stage 3
in4a = get_inception_a(data, 224, 64, 96, 96, 128, "avg", 128, '4a')
in4b = get_inception_a(in4a, 192, 96, 128, 96, 128, "avg", 128, '4b')
in4c = get_inception_a(in4b, 160, 128, 160, 128, 160, "avg", 128, '4c')
in4d = get_inception_a(in4c, 96, 128, 192, 160, 192, "avg", 128, '4d')
in4e = get_inception_b(in4d, 128, 192, 192, 256, '4e')
# stage 4
in5a = get_inception_a(data, 352, 192, 320, 160, 224, "avg", 128, '5a')
in5b = get_inception_a(in5a, 352, 192, 320, 192, 224, "max", 128, '5b')
# stage 5
pool = mx.symbol.Pooling(data=data,
kernel=(7, 7),
stride=(1, 1),
name="global_pool",
pool_type='avg')
flatten = mx.symbol.Flatten(data=pool)
self._stage_1 = SymbolBlock(outputs=conv2, inputs=mx.sym.var('data'))
self._stage_2 = SymbolBlock(outputs=in3c, inputs=mx.sym.var('data'))
self._stage_3 = SymbolBlock(outputs=in4e, inputs=mx.sym.var('data'))
self._stage_4 = SymbolBlock(outputs=in5b, inputs=mx.sym.var('data'))
self._stage_5 = SymbolBlock(outputs=flatten, inputs=mx.sym.var('data'))
class InceptionBNStages(HybridBlock):
def __init__(self, **kwargs):
super(InceptionBNStages, self).__init__(**kwargs)
# data
data = mx.symbol.Variable(name="data")
# stage 1
conv1 = get_conv(data=data,
num_filter=64,
kernel=(7, 7),
stride=(2, 2),
pad=(3, 3),
name='1')
pool1 = mx.symbol.Pooling(data=conv1,
kernel=(3, 3),
stride=(2, 2),
name='pool_1',
pool_type='max')
conv2red = get_conv(data=pool1,
num_filter=64,
kernel=(1, 1),
stride=(1, 1),
name='2_red')
conv2 = get_conv(data=conv2red,
num_filter=192,
kernel=(3, 3),
stride=(1, 1),
pad=(1, 1),
name='2')
# stage 2
pool2 = mx.symbol.Pooling(data=data,
kernel=(3, 3),
stride=(2, 2),
name='pool_2',
pool_type='max')
in3a = get_inception_a(pool2, 64, 64, 64, 64, 96, "avg", 32, '3a')
in3b = get_inception_a(in3a, 64, 64, 96, 64, 96, "avg", 64, '3b')
in3c = get_inception_b(in3b, 128, 160, 64, 96, '3c')
# stage 3
in4a = get_inception_a(data, 224, 64, 96, 96, 128, "avg", 128, '4a')
in4b = get_inception_a(in4a, 192, 96, 128, 96, 128, "avg", 128, '4b')
in4c = get_inception_a(in4b, 160, 128, 160, 128, 160, "avg", 128, '4c')
in4d = get_inception_a(in4c, 96, 128, 192, 160, 192, "avg", 128, '4d')
in4e = get_inception_b(in4d, 128, 192, 192, 256, '4e')
# stage 4
in5a = get_inception_a(data, 352, 192, 320, 160, 224, "avg", 128, '5a')
in5b = get_inception_a(in5a, 352, 192, 320, 192, 224, "max", 128, '5b')
# stage 5
pool = mx.symbol.Pooling(data=data,
kernel=(7, 7),
stride=(1, 1),
name="global_pool",
pool_type='avg')
flatten = mx.symbol.Flatten(data=pool)
self._stage_1 = SymbolBlock(outputs=conv2, inputs=mx.sym.var('data'))
self._stage_2 = SymbolBlock(outputs=in3c, inputs=mx.sym.var('data'))
self._stage_3 = SymbolBlock(outputs=in4e, inputs=mx.sym.var('data'))
self._stage_4 = SymbolBlock(outputs=in5b, inputs=mx.sym.var('data'))
self._stage_5 = SymbolBlock(outputs=flatten, inputs=mx.sym.var('data'))
def hybrid_forward(self, F, x):
stage1 = self._stage_1(x)
stage2 = self._stage_2(stage1)
stage3 = self._stage_3(stage2)
stage4 = self._stage_4(stage3)
stage5 = self._stage_5(stage4)
return stage1, stage2, stage3, stage4, stage5
def load(self, param_file, ctx=mx.cpu()):
self._stage_1.collect_params().load(param_file,
ctx=ctx,
ignore_extra=True)
self._stage_2.collect_params().load(param_file,
ctx=ctx,
ignore_extra=True)
self._stage_3.collect_params().load(param_file,
ctx=ctx,
ignore_extra=True)
self._stage_4.collect_params().load(param_file,
ctx=ctx,
ignore_extra=True)
self._stage_5.collect_params().load(param_file,
ctx=ctx,
ignore_extra=True)