From 5c81473ff5804c086cdd357cbbd8d73af977cb1e Mon Sep 17 00:00:00 2001 From: reminisce Date: Tue, 16 Jan 2018 11:54:04 -0800 Subject: [PATCH] Fix unit tests --- python/mxnet/quantization.py | 2 +- tests/python/unittest/test_quantization.py | 6 ++++-- 2 files changed, 5 insertions(+), 3 deletions(-) diff --git a/python/mxnet/quantization.py b/python/mxnet/quantization.py index 2ff926341f91..e1531161f44e 100644 --- a/python/mxnet/quantization.py +++ b/python/mxnet/quantization.py @@ -205,7 +205,7 @@ def _get_optimal_threshold(arr, num_bins=8001, num_quantized_bins=255): hist, hist_edeges = np.histogram(arr, bins=num_bins, range=(-th, th)) zero_bin_idx = num_bins / 2 num_half_quantized_bins = num_quantized_bins / 2 - assert np.allclose(hist_edeges[zero_bin_idx] + hist_edeges[zero_bin_idx + 1], 0, rtol=1e-5, atol=1e-7) + assert np.allclose(hist_edeges[int(zero_bin_idx)] + hist_edeges[int(zero_bin_idx + 1)], 0, rtol=1e-5, atol=1e-7) thresholds = np.zeros(num_bins / 2 + 1 - num_quantized_bins / 2) divergence = np.zeros_like(thresholds) diff --git a/tests/python/unittest/test_quantization.py b/tests/python/unittest/test_quantization.py index d21e2a6d2803..8f2df9b5872d 100644 --- a/tests/python/unittest/test_quantization.py +++ b/tests/python/unittest/test_quantization.py @@ -1,7 +1,9 @@ """Some of the tests using CUDNN require a special GPU instruction called dp4a. Ref: http://images.nvidia.com/content/pdf/tesla/184457-Tesla-P4-Datasheet-NV-Final-Letter-Web.pdf """ -from mxnet.test_utils import * +import mxnet as mx +import numpy as np +from mxnet.test_utils import assert_almost_equal, rand_ndarray, rand_shape_nd, same def test_quantize_float32_to_int8(): @@ -326,7 +328,7 @@ def test_quantize_sym_with_calib(): assert_almost_equal(np.array([lhs]), np.array([rhs])) lhs = float(attr_dict[name]['max_calib_range']) rhs = th_dict[op_name_to_th_name[name]][1] - assert_almost_equal(np.array([lhs]), np.array([rhs])) + assert_almost_equal(np.array([lhs]), np.array([rhs]), rtol=1e-3, atol=1e-5) def test_get_optimal_thresholds():