You can not select more than 25 topics Topics must start with a chinese character,a letter or number, can include dashes ('-') and can be up to 35 characters long.

test_fake_quant.py 4.9 kB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152
  1. # -*- coding: utf-8 -*-
  2. # MegEngine is Licensed under the Apache License, Version 2.0 (the "License")
  3. #
  4. # Copyright (c) 2014-2021 Megvii Inc. All rights reserved.
  5. #
  6. # Unless required by applicable law or agreed to in writing,
  7. # software distributed under the License is distributed on an
  8. # "AS IS" BASIS, WITHOUT ARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  9. import numpy as np
  10. import pytest
  11. import megengine as mge
  12. from megengine import tensor
  13. from megengine.core.autodiff.grad import Function, Grad
  14. from megengine.core.tensor.dtype import QuantDtypeMeta
  15. from megengine.core.tensor.utils import make_shape_tuple
  16. from megengine.quantization.internal_fake_quant import *
  17. from megengine.quantization.utils import (
  18. QuantMode,
  19. create_qparams,
  20. fake_quant_tensor,
  21. tqt_forward,
  22. )
  23. class TQT_numpy:
  24. def __init__(self, lowerbound, upperbound):
  25. super().__init__()
  26. self.lowerbound = lowerbound
  27. self.upperbound = upperbound
  28. def forward(self, inp, scale):
  29. t = 2 ** scale
  30. # t = F.maximum(t, 1e-4)
  31. inp_scaled = inp / t
  32. inp_clipped = np.maximum(
  33. np.minimum(inp_scaled, self.upperbound), self.lowerbound
  34. )
  35. inp_rounded = np.round(inp_clipped)
  36. inp_flq = inp_rounded * t
  37. self.saved_tensors = (inp_scaled, inp_rounded, t)
  38. return inp_flq
  39. def backward(self, grad_inp_flq):
  40. (inp_scaled, inp_rounded, t) = self.saved_tensors
  41. mask_clip = (inp_scaled < -0.5 + self.lowerbound) + (
  42. inp_scaled > self.upperbound + 0.5
  43. ) # mask for accumulating the gradients of |data_scaled|>L
  44. mask_quant = np.abs(
  45. mask_clip - 1
  46. ) # mask for accumulating the gradients with |data_scaled|<=L
  47. grad_quant = (
  48. grad_inp_flq * mask_quant * (inp_rounded - inp_scaled)
  49. ) # gradient within |data_scaled|<=L
  50. grad_clip = (
  51. grad_inp_flq * mask_clip * inp_rounded
  52. ) # gradient with | data_scaled|>L
  53. grad_s = grad_clip.sum() + grad_quant.sum()
  54. # dL/ds = dL/dt * t * ln(2)
  55. grad_s = grad_s * t * np.log(2)
  56. grad_inp = grad_inp_flq * mask_quant
  57. return grad_inp, grad_s
  58. def test_tqt():
  59. g = []
  60. def cb(grad):
  61. g.append(grad)
  62. x = np.random.randint(-128, 128, size=(1, 2, 3, 4)).astype("float32")
  63. s = np.random.rand(1) - 1
  64. g_y = np.ones(shape=(1, 2, 3, 4), dtype="float32")
  65. n = TQT_numpy(-127, 127)
  66. y_np = n.forward(x, s)
  67. g_x_np, g_s_np = n.backward(g_y)
  68. x = mge.tensor(x, dtype="float32")
  69. s = mge.tensor(s, dtype="float32")
  70. g_y = mge.tensor(g_y, dtype="float32")
  71. grad = Grad().wrt(x, s, callback=cb)
  72. y = tqt_forward(-127, 127, x, s)
  73. grad(y, g_y)
  74. g_x, g_s = g
  75. np.testing.assert_allclose(y.numpy(), y_np, rtol=1e-5, atol=1e-5)
  76. np.testing.assert_allclose(g_x.numpy(), g_x_np, rtol=1e-5, atol=1e-5)
  77. np.testing.assert_allclose(g_s.numpy(), g_s_np, rtol=5e-5, atol=5e-5)
  78. def _save_to(self, name="grad"):
  79. def callback(grad):
  80. setattr(self, name, grad)
  81. return callback
  82. class Round(Function):
  83. def forward(self, x):
  84. return F.round(x)
  85. def backward(self, output_grads):
  86. return output_grads
  87. def fake_quant_tensor_gt(inp, scale, zero_point, qmin, qmax):
  88. oup = Round()(inp / scale) + zero_point
  89. oup = F.minimum(F.maximum(oup, qmin), qmax)
  90. oup = (oup - zero_point) * scale
  91. return oup
  92. def test_fakequant():
  93. qmin = -126
  94. qmax = 129
  95. test_dtype = QuantDtypeMeta("test_qint8", None, "int8", qmin, qmax)
  96. def run(zero_point, scale):
  97. qparams = create_qparams(QuantMode.ASYMMERTIC, test_dtype, scale, zero_point)
  98. inp_data = np.random.uniform(low=-512.0, high=512.0, size=(1, 32, 32, 32))
  99. inp = tensor(inp_data, dtype=np.float32)
  100. # test forward
  101. oup = fake_quant_tensor(inp, qparams).numpy()
  102. oup_gt = fake_quant_tensor_gt(inp, scale, zero_point, qmin, qmax).numpy()
  103. assert np.allclose(oup, oup_gt)
  104. assert oup.shape == oup_gt.shape
  105. # test backward
  106. x = tensor(inp_data, dtype=np.float32)
  107. grad = Grad().wrt(x, callback=_save_to(x))
  108. y = fake_quant_tensor(x, qparams)
  109. grad(y, tensor(F.ones_like(x)))
  110. x1 = tensor(inp_data, dtype=np.float32)
  111. grad = Grad().wrt(x1, callback=_save_to(x1))
  112. y1 = fake_quant_tensor_gt(x1, scale, zero_point, qmin, qmax)
  113. grad(y1, tensor(F.ones_like(x1)))
  114. assert np.allclose(x.grad.numpy(), x1.grad.numpy())
  115. assert make_shape_tuple(x.grad.shape) == make_shape_tuple(x1.grad.shape)
  116. zero_point = tensor([1.0], dtype=np.float32)
  117. scale = tensor([4.0], dtype=np.float32)
  118. run(zero_point, scale)
  119. zero_point = tensor(1.0 * np.ones((1, 32, 1, 1)), dtype=np.float32)
  120. scale = tensor(4.0 * np.ones((1, 32, 1, 1)), dtype=np.float32)
  121. run(zero_point, scale)

MegEngine 安装包中集成了使用 GPU 运行代码所需的 CUDA 环境,不用区分 CPU 和 GPU 版。 如果想要运行 GPU 程序,请确保机器本身配有 GPU 硬件设备并安装好驱动。 如果你想体验在云端 GPU 算力平台进行深度学习开发的感觉,欢迎访问 MegStudio 平台