| # 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. |
| |
| import mxnet as mx |
| from mxnet import np, npx |
| from mxnet.test_utils import same, use_np, assert_almost_equal |
| import random |
| import pytest |
| |
| @use_np |
| @pytest.mark.parametrize('shape', |
| [(3, 2), (9,17), (2, 7, 1, 8)] + [(i,) for i in range(1,65)]) |
| def test_contrib_intgemm_maxabsolute(shape): |
| if "intgemm_maxabsolute" not in dir(mx.nd.contrib): |
| return |
| # mx.nd API |
| m = mx.nd.random_uniform(low=-100.0, high=100.0, shape=shape) |
| fast = mx.nd.contrib.intgemm_maxabsolute(m) |
| slow = mx.nd.max(mx.nd.abs(m)) |
| assert same(fast, slow) |
| # np API |
| m = np.random.uniform(low=-100.0, high=100.0, size=shape) |
| fast = npx.intgemm_maxabsolute(m).reshape(()) |
| slow = np.max(np.abs(m)) |
| assert same(fast, slow) |
| |
| @use_np |
| @pytest.mark.parametrize('shape', [(i,) for i in range(1, 67)] + [(2,3), (130, 12)]) |
| @pytest.mark.parametrize('max_quant', [2.0]) |
| def test_contrib_intgemm_prepare_data(shape, max_quant): |
| if "intgemm_prepare_data" not in dir(mx.nd.contrib): |
| return |
| m = mx.nd.random_uniform(low=-3.0, high=3.0, shape=shape) |
| scaled = m * 127.0 / max_quant |
| # Rounding 0.5 can go up or down. Move values away from 0.5. |
| too_close = mx.nd.abs(mx.nd.round(scaled) - scaled) > 0.45 |
| # Add 0.2 in scaled space so (0.45, 0.55) maps to (0.65, 0.75) which will round consistently. |
| m += max_quant / 127.0 * 0.2 * too_close |
| |
| # Reference: scale and round |
| ref = mx.nd.round(m * 127.0 / max_quant) |
| # Clip to [-127, 127]. Because otherwise e.g. -129 casts to +127. |
| ref = mx.nd.broadcast_maximum(ref, mx.nd.array([-127.0])) |
| ref = mx.nd.broadcast_minimum(ref, mx.nd.array([127.0])) |
| # Reference: cast to int8 |
| ref = mx.nd.cast(ref, dtype='int8') |
| # Reference: ban -128 |
| ref = mx.nd.broadcast_maximum(ref, mx.nd.array([-127], dtype = 'int8')) |
| |
| test = mx.nd.contrib.intgemm_prepare_data(m, mx.nd.array([max_quant])) |
| assert same(test, ref) |
| test = npx.intgemm_prepare_data(m.as_np_ndarray(), np.array([max_quant])) |
| assert same(test, ref.as_np_ndarray()) |
| |
| @use_np |
| @pytest.mark.parametrize('shape', [(8, 64), (16, 64), (8, 128), (16, 128), (2, 4, 64)]) |
| @pytest.mark.parametrize('max_quant', [0.2, 3.0]) |
| @pytest.mark.parametrize('api', [(mx.nd.contrib, mx.nd), (npx, np)]) |
| def test_contrib_intgemm_weight_consistent(shape, max_quant, api): |
| # The weight format is actually CPU-dependent so we don't directly test the |
| # output, but indirectly test that it works. |
| if "intgemm_prepare_weight" not in dir(mx.nd.contrib): |
| return |
| contrib, top = api |
| max_array = top.array([max_quant]) |
| if top == mx.nd: |
| m = top.random_uniform(low=-3.0, high=3.0, shape=shape) |
| else: |
| m = np.random.uniform(size=shape) |
| direct = contrib.intgemm_prepare_weight(m, max_array) |
| quant = contrib.intgemm_prepare_data(m, max_array) |
| indirect = contrib.intgemm_prepare_weight(quant, already_quantized=True) |
| # Should get the same data from direct call and already_quantized version. |
| assert same(direct, indirect) |
| |
| @use_np |
| @pytest.mark.parametrize('indices', [ |
| [0,1,2,3,4,5,6,7], |
| [1,2,1,2,1,2,1,2], |
| [7,6,5,4,3,2,1,0], |
| [3,1,4,1,5,9,2,6], |
| # Since random_uniform doesn't support int8, use python |
| [random.randint(0,15) for i in range(8)], |
| [random.randint(0,15) for i in range(16)], |
| [random.randint(0,15) for i in range(24)] |
| ]) |
| @pytest.mark.parametrize('api', [(mx.nd.contrib, mx.nd), (npx, np)]) |
| def test_contrib_intgemm_take_weight(indices, api): |
| if "intgemm_take_weight" not in dir(mx.nd.contrib): |
| return |
| contrib, top = api |
| m = top.array([random.randint(-127,127) for i in range(16 * 64)], dtype='int8') |
| m = m.reshape((16, 64)) |
| indices = top.array(indices, dtype='int32') |
| # Prepare weight then take. |
| test = contrib.intgemm_prepare_weight(m, already_quantized=True) |
| test = contrib.intgemm_take_weight(test, indices) |
| # Take then prepare. |
| ref = m.take(indices, axis=0) |
| ref = contrib.intgemm_prepare_weight(ref, already_quantized=True) |
| assert same(test, ref) |
| |
| @use_np |
| @pytest.mark.parametrize('data_rows', range(1, 5)) |
| @pytest.mark.parametrize('inner', range(64, 256, 64)) |
| @pytest.mark.parametrize('weight_cols', range(8, 24, 8)) |
| @pytest.mark.parametrize('api', [ |
| (mx.nd.contrib, mx.nd, mx.nd.FullyConnected, mx.nd.cast), |
| (npx, np, npx.fully_connected, npx.cast)]) |
| def test_contrib_intgemm_multiply(data_rows, inner, weight_cols, api): |
| if "intgemm_fully_connected" not in dir(mx.nd.contrib): |
| return |
| contrib, top, fully_connected, cast = api |
| #The multiplication routine has approximations so everything is tested |
| #deterministically to ensure bounds are met. |
| random.seed(1) |
| |
| # Don't use full range (-127, 127) to avoid saturation. |
| data = [random.randint(-64, 64) for i in range(data_rows * inner)] |
| data = top.array(data, dtype='int8').reshape((data_rows, inner)) |
| weight = [random.randint(-64, 64) for i in range(inner * weight_cols)] |
| weight = top.array(weight, dtype='int8').reshape((weight_cols, inner)) |
| weight_prepared = contrib.intgemm_prepare_weight(weight, already_quantized=True) |
| |
| # int32 output, no bias |
| test = contrib.intgemm_fully_connected(data, |
| weight_prepared, |
| no_bias=True, |
| flatten=False, |
| out_type='int32', |
| num_hidden=weight_cols) |
| ref = fully_connected(cast(data, dtype='float32'), |
| cast(weight, dtype='float32'), |
| no_bias=True, |
| flatten=False, |
| num_hidden=weight_cols) |
| assert_almost_equal(cast(test, dtype='float32'), ref, rtol=0.01, atol=0.01) |
| |
| # float32 output, no bias |
| scale = 3.0 |
| test = contrib.intgemm_fully_connected(data, |
| weight_prepared, |
| top.array([scale]), |
| no_bias=True, |
| flatten=False, |
| out_type='float32', |
| num_hidden=weight_cols) |
| assert_almost_equal(test, ref * scale, rtol=0.01, atol=0.01) |
| |
| # int32 output, bias |
| bias = top.array([random.randint(-60000, 60000) for i in range(weight_cols)], dtype = 'int32') |
| test = contrib.intgemm_fully_connected(data, |
| weight_prepared, |
| bias, |
| no_bias=False, |
| flatten=False, |
| out_type='int32', |
| num_hidden=weight_cols) |
| ref = fully_connected(cast(data, dtype='float32'), |
| cast(weight, dtype='float32'), |
| cast(bias, dtype='float32'), |
| no_bias=False, |
| flatten=False, |
| num_hidden=weight_cols) |
| assert_almost_equal(cast(test, dtype='float32'), ref, rtol=0.01, atol=0.01) |
| |
| # float32 output, bias |
| # Scaling is applied before bias (and bias is not scaled). So to make the |
| # reference comparison easy, just scale the bias beforehand. |
| test = contrib.intgemm_fully_connected(data, |
| weight_prepared, |
| top.array([scale]), |
| cast(bias, dtype='float32') * scale, |
| no_bias=False, |
| flatten=False, |
| out_type='float32', |
| num_hidden=weight_cols) |
| assert_almost_equal(test, ref * scale, rtol=0.01, atol=0.01) |
| |
| # float32 input should work the same as manually prepared int8 input. |
| data_float = top.array([random.uniform(-3.14, 3.14) for i in range(data_rows * inner)]) |
| data_float = data_float.reshape(data_rows, inner) |
| direct = contrib.intgemm_fully_connected(data_float, |
| weight_prepared, |
| top.array([scale]), |
| cast(bias, dtype='float32'), |
| no_bias=False, |
| flatten=False, |
| out_type='float32', |
| num_hidden=weight_cols) |
| maxabs = contrib.intgemm_maxabsolute(data_float) |
| data_prepared = contrib.intgemm_prepare_data(data_float, maxabs) |
| cooked = contrib.intgemm_fully_connected(data_prepared, |
| weight_prepared, |
| top.array(scale * maxabs / 127.0), |
| cast(bias, dtype='float32'), |
| no_bias=False, |
| flatten=False, |
| out_type='float32', |
| num_hidden=weight_cols) |
| assert_almost_equal(direct, cooked, rtol=0.01, atol=0.01) |
