| # Licensed to the Apache Software Foundation (ASF) under one |
| # or more contributor license agreements. See the NOTICE file |
| # distributed with this work for additional information |
| # regarding copyright ownership. The ASF licenses this file |
| # to you 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. |
| |
| # pylint: skip-file |
| from __future__ import print_function |
| import numpy as np |
| import mxnet as mx |
| import copy |
| from mxnet.test_utils import * |
| import pytest |
| from mxnet.gluon.model_zoo.vision import get_model |
| |
| def make_subgraph(subg, *args): |
| js = subg.tojson() |
| return subg |
| |
| @pytest.mark.serial |
| def test_make_subgraph(): |
| def make_subgraph1(stype): |
| a = mx.symbol.Variable(name='a', stype=stype) |
| b = mx.symbol.Variable(name='b', stype=stype) |
| c = a * b |
| d = c * 2 |
| |
| a1 = mx.symbol.Variable(name='a', stype=stype) |
| b1 = mx.symbol.Variable(name='b', stype=stype) |
| y = make_subgraph(c, a1, b1) |
| y = y * 2 |
| |
| s = (10, 10) |
| a_arr = mx.nd.array(np.random.normal(-0.1, 0.1, size=s), |
| ctx=default_device()).tostype(stype) |
| b_arr = mx.nd.array(np.random.normal(-0.1, 0.1, size=s), |
| ctx=default_device()).tostype(stype) |
| return (d, y, {'a': a_arr, 'b': b_arr}, {}) |
| |
| def create_weights(shapes, names): |
| nd_dict = {} |
| sym_dict = {} |
| assert len(shapes) == len(names) |
| for i in range(len(shapes)): |
| sym_dict[names[i]] = mx.symbol.Variable(names[i]) |
| nd_dict[names[i]] = mx.nd.array(np.ones(shapes[i]), ctx=default_device()) |
| return (nd_dict, sym_dict) |
| |
| def make_subgraph_weight(orig, shape, stype): |
| arg_shapes, out_shapes, aux_shapes = orig.infer_shape(data=shape) |
| weight_shapes = arg_shapes[1:] |
| weight_names = orig.list_arguments()[1:] |
| weight_dict, weight_sym_dict = create_weights(weight_shapes, weight_names) |
| aux_dict, aux_sym_dict = create_weights(aux_shapes, orig.list_auxiliary_states()) |
| |
| input_dict = copy.deepcopy(weight_sym_dict) |
| input_dict.update(aux_sym_dict) |
| input_dict['data'] = mx.symbol.Variable('data', stype=stype) |
| input_list = [] |
| for name in orig.list_inputs(): |
| assert name in input_dict.keys() |
| input_list.append(input_dict[name]) |
| subg = make_subgraph(orig, *input_list) |
| |
| arr = mx.nd.random.uniform(-1, 1, shape=shape, ctx=default_device()).tostype(stype) |
| arg_dict = weight_dict |
| arg_dict['data'] = arr |
| return (orig, subg, arg_dict, aux_dict) |
| |
| def make_subgraph2(stype, out_mean_var): |
| data = mx.symbol.Variable('data', stype=stype) |
| orig = mx.symbol.BatchNorm(data, fix_gamma=False, |
| output_mean_var=out_mean_var, name="batchnorm") |
| s = (10, 10) |
| return make_subgraph_weight(orig, s, stype) |
| |
| def make_subgraph3(stype): |
| data = mx.symbol.Variable('data', stype=stype) |
| conv1 = mx.symbol.Convolution(data=data, kernel=(3, 3), num_filter=16, no_bias=True) |
| bn1 = mx.symbol.BatchNorm(conv1, fix_gamma=False, output_mean_var=False) |
| conv2 = mx.symbol.Convolution(data=data, kernel=(3, 3), num_filter=16, no_bias=True) |
| bn2 = mx.symbol.BatchNorm(conv2, fix_gamma=False, output_mean_var=False) |
| orig = bn1 + bn2 |
| s = (1, 3, 32, 32) |
| return make_subgraph_weight(orig, s, stype) |
| |
| def make_subgraph4(stype): |
| model = get_model('resnet18_v1') |
| model.hybridize() |
| model.initialize() |
| s = (1, 3, 32, 32) |
| data = mx.np.random.normal(size=s) |
| out = model(data) |
| model.export('resnet18') |
| orig = mx.sym.load('resnet18-symbol.json') |
| return make_subgraph_weight(orig, s, stype) |
| |
| make_subgraphs = [make_subgraph1, |
| lambda stype: make_subgraph2(stype, False), |
| lambda stype: make_subgraph2(stype, True), |
| make_subgraph3, make_subgraph4] |
| stypes = ['default', 'row_sparse'] |
| for make_subg in make_subgraphs: |
| for stype in stypes: |
| orig, subg, inputs, aux_states = make_subg(stype) |
| all_inputs = copy.deepcopy(inputs) |
| all_inputs.update(aux_states) |
| args_grad = {key : mx.nd.empty(shape=all_inputs[key].shape) for key in all_inputs.keys()} |
| e1 = orig._bind(ctx=default_device(), args=all_inputs, args_grad=args_grad, |
| aux_states=all_inputs) |
| args_grad = {key : mx.nd.empty(shape=all_inputs[key].shape) for key in all_inputs.keys()} |
| e2 = subg._bind(ctx=default_device(), args=all_inputs, args_grad=args_grad, |
| aux_states=all_inputs) |
| e1.forward(is_train=True) |
| e2.forward(is_train=True) |
| for i in range(len(e1.outputs)): |
| assert_almost_equal(e1.outputs[i].asnumpy(), e2.outputs[i].asnumpy(), |
| rtol=0.001, atol=0.0001) |
| |
| out_grads = [mx.nd.random.uniform(-1, 1, shape=out.shape, ctx=default_device()) |
| for out in e1.outputs] |
| e1.backward(out_grads) |
| e2.backward(out_grads) |
| for i in range(len(e1.grad_arrays)): |
| assert_almost_equal(e1.grad_arrays[i].asnumpy(), e2.grad_arrays[i].asnumpy(), |
| rtol=0.001, atol=0.0001) |
| |
| |
| @pytest.mark.serial |
| def test_subgraph_with_customOp(): |
| class MyAdd(mx.operator.CustomOp): |
| def forward(self, is_train, req, in_data, out_data, aux): |
| self.assign(out_data[0], req[0], in_data[0] + 1) |
| |
| def backward(self, req, out_grad, in_data, out_data, in_grad, aux): |
| self.assign(in_grad[0], req[0], out_grad[0]) |
| |
| @mx.operator.register('MyAdd1') |
| class MyAdd1Prop(mx.operator.CustomOpProp): |
| def __init__(self): |
| super(MyAdd1Prop, self).__init__(need_top_grad=True) |
| |
| def list_arguments(self): |
| return ['data'] |
| |
| def list_outputs(self): |
| return ['output'] |
| |
| def infer_shape(self, in_shape): |
| # inputs, outputs, aux |
| return [in_shape[0]], [in_shape[0]], [] |
| |
| def create_operator(self, ctx, shapes, dtypes): |
| return MyAdd() |
| |
| @mx.operator.register('MyAdd2') |
| class MyAdd2Prop(mx.operator.CustomOpProp): |
| def __init__(self): |
| super(MyAdd2Prop, self).__init__(need_top_grad=True) |
| |
| def list_arguments(self): |
| return ['data'] |
| |
| def list_outputs(self): |
| return ['output'] |
| |
| def infer_shape(self, in_shape): |
| # inputs, outputs, aux |
| return [in_shape[0]], [in_shape[0]], [] |
| |
| def create_operator(self, ctx, shapes, dtypes): |
| return MyAdd() |
| |
| inp = mx.nd.zeros(shape=(100, 100)) |
| a = mx.symbol.Variable('a') |
| b = a + 1 |
| b = mx.symbol.Custom(data=a, op_type='MyAdd1') |
| c = mx.symbol.Custom(data=a, op_type='MyAdd2') |
| b._bind(mx.cpu(), {'a': inp}).forward() |
| c._bind(mx.cpu(), {'a': inp}).forward() |
| mx.nd.waitall() |
