Skip to content
体验新版
项目
组织
正在加载...
登录
切换导航
打开侧边栏
MegEngine 天元
MegEngine
提交
2694ff81
MegEngine
项目概览
MegEngine 天元
/
MegEngine
1 年多 前同步成功
通知
403
Star
4705
Fork
582
代码
文件
提交
分支
Tags
贡献者
分支图
Diff
Issue
0
列表
看板
标记
里程碑
合并请求
0
DevOps
流水线
流水线任务
计划
Wiki
0
Wiki
分析
仓库
DevOps
项目成员
Pages
MegEngine
项目概览
项目概览
详情
发布
仓库
仓库
文件
提交
分支
标签
贡献者
分支图
比较
Issue
0
Issue
0
列表
看板
标记
里程碑
合并请求
0
合并请求
0
Pages
DevOps
DevOps
流水线
流水线任务
计划
分析
分析
仓库分析
DevOps
Wiki
0
Wiki
成员
成员
收起侧边栏
关闭侧边栏
动态
分支图
创建新Issue
流水线任务
提交
Issue看板
提交
2694ff81
编写于
10月 25, 2022
作者:
M
Megvii Engine Team
浏览文件
操作
浏览文件
下载
电子邮件补丁
差异文件
fix(mge/module): fix some deconv fuse bn problem
GitOrigin-RevId: e88a63328065234c8060317d7c7e558c95bfbd4b
上级
1886ebc3
变更
4
隐藏空白更改
内联
并排
Showing
4 changed file
with
199 addition
and
111 deletion
+199
-111
imperative/python/megengine/module/qat/conv_transpose_bn.py
imperative/python/megengine/module/qat/conv_transpose_bn.py
+3
-4
imperative/python/megengine/utils/bn_fusion.py
imperative/python/megengine/utils/bn_fusion.py
+39
-60
imperative/python/test/unit/module/test_qat.py
imperative/python/test/unit/module/test_qat.py
+65
-47
imperative/python/test/unit/quantization/test_quantize.py
imperative/python/test/unit/quantization/test_quantize.py
+92
-0
未找到文件。
imperative/python/megengine/module/qat/conv_transpose_bn.py
浏览文件 @
2694ff81
...
...
@@ -24,8 +24,7 @@ class _ConvTransposeBnActivation2d(Float._ConvTransposeBnActivation2d, QATModule
# get fold bn conv_transpose2d param
gamma
=
self
.
bn
.
weight
if
gamma
is
None
:
gamma
=
ones
((
self
.
bn
.
num_features
),
dtype
=
"float32"
)
gamma
=
gamma
.
reshape
(
1
,
-
1
,
1
,
1
)
gamma
=
ones
((
1
,
self
.
bn
.
num_features
,
1
,
1
),
dtype
=
"float32"
)
beta
=
self
.
bn
.
bias
if
beta
is
None
:
beta
=
zeros
((
1
,
self
.
bn
.
num_features
,
1
,
1
),
dtype
=
"float32"
)
...
...
@@ -44,10 +43,10 @@ class _ConvTransposeBnActivation2d(Float._ConvTransposeBnActivation2d, QATModule
bn_istd
=
1.0
/
sqrt
(
bn_var
+
self
.
bn
.
eps
)
scale_factor
=
gamma
*
bn_istd
if
self
.
conv_transpose2d
.
groups
==
1
:
w_fold
=
self
.
conv_transpose2d
.
weight
*
scale_factor
.
reshape
(
-
1
,
1
,
1
,
1
)
w_fold
=
self
.
conv_transpose2d
.
weight
*
scale_factor
.
reshape
(
1
,
-
1
,
1
,
1
)
else
:
w_fold
=
self
.
conv_transpose2d
.
weight
*
scale_factor
.
reshape
(
self
.
conv_transpose2d
.
groups
,
-
1
,
1
,
1
,
1
self
.
conv_transpose2d
.
groups
,
1
,
-
1
,
1
,
1
)
w_fold
=
self
.
apply_quant_weight
(
w_fold
)
...
...
imperative/python/megengine/utils/bn_fusion.py
浏览文件 @
2694ff81
...
...
@@ -32,15 +32,21 @@ _MAP_TO_FUSED_MODULE = {
def
fold_weight_bias
(
weight
,
bias
,
gamma
,
beta
,
bn_mean
,
bn_var
,
eps
=
1e-5
,
transpose
=
False
):
shape
=
(
1
,
-
1
,
1
,
1
)
shape
=
(
-
1
,
1
,
1
,
1
)
if
transpose
:
shape
=
(
-
1
,
1
,
1
,
1
)
shape
=
(
1
,
-
1
,
1
,
1
)
kernel_shape
=
weight
.
shape
if
len
(
kernel_shape
)
==
5
:
groups
,
num_features
=
kernel_shape
[
0
],
kernel_shape
[
1
]
if
transpose
:
groups
,
num_features
=
kernel_shape
[
0
],
kernel_shape
[
2
]
else
:
groups
,
num_features
=
kernel_shape
[
0
],
kernel_shape
[
1
]
else
:
groups
,
num_features
=
1
,
kernel_shape
[
0
]
if
transpose
:
groups
,
num_features
=
1
,
kernel_shape
[
1
]
else
:
groups
,
num_features
=
1
,
kernel_shape
[
0
]
out_channels
=
groups
*
num_features
if
gamma
is
None
:
...
...
@@ -93,12 +99,37 @@ def fuse_conv_bn_relu_module(conv: Conv2d, bn: BatchNorm2d, relu: ReLU):
compute_mode
=
conv
.
compute_mode
,
name
=
conv
.
name
,
)
new_conv
=
module
if
bn
is
None
or
not
conv
.
training
else
module
.
conv
if
isinstance
(
conv
,
ConvTranspose2d
):
module
.
output_padding
=
conv
.
output_padding
new_conv
=
(
module
if
bn
is
None
or
not
conv
.
training
else
module
.
conv_transpose2d
)
else
:
new_conv
=
module
if
bn
is
None
or
not
conv
.
training
else
module
.
conv
weight
,
bias
=
conv
.
weight
,
conv
.
bias
if
not
conv
.
training
and
bn
is
not
None
:
weight
,
bias
=
fold_weight_bias
(
weight
,
bias
,
bn
.
weight
,
bn
.
bias
,
bn
.
running_mean
,
bn
.
running_var
,
bn
.
eps
,
)
if
isinstance
(
conv
,
ConvTranspose2d
):
weight
,
bias
=
fold_weight_bias
(
weight
,
bias
,
bn
.
weight
,
bn
.
bias
,
bn
.
running_mean
,
bn
.
running_var
,
bn
.
eps
,
transpose
=
True
,
)
else
:
weight
,
bias
=
fold_weight_bias
(
weight
,
bias
,
bn
.
weight
,
bn
.
bias
,
bn
.
running_mean
,
bn
.
running_var
,
bn
.
eps
,
)
new_conv
.
weight
=
Parameter
(
weight
)
if
bias
is
not
None
:
new_conv
.
bias
=
Parameter
(
bias
)
...
...
@@ -106,55 +137,3 @@ def fuse_conv_bn_relu_module(conv: Conv2d, bn: BatchNorm2d, relu: ReLU):
module
.
bn
=
deepcopy
(
bn
)
new_conv
.
training
=
conv
.
training
return
module
def
fuse_conv_transpose2d_bn_relu_module
(
conv_transpose2d
:
ConvTranspose2d
,
bn
:
BatchNorm2d
,
relu
:
ReLU
):
module_key
=
tuple
([
type
(
m
)
for
m
in
[
conv_transpose2d
,
bn
,
relu
]
if
m
])
if
bn
:
assert
(
conv_transpose2d
.
training
==
bn
.
training
),
"ConvTranspose2d and BN both must be in the same mode (train or eval)."
assert
(
bn
.
num_features
==
conv_transpose2d
.
out_channels
),
"Output channel of ConvTranspose2d must match num_features of BatchNorm2d"
module_key
=
module_key
+
(
conv_transpose2d
.
training
,)
module
=
_MAP_TO_FUSED_MODULE
[
module_key
](
in_channels
=
conv_transpose2d
.
in_channels
,
out_channels
=
conv_transpose2d
.
out_channels
,
kernel_size
=
conv_transpose2d
.
kernel_size
,
stride
=
conv_transpose2d
.
stride
,
padding
=
conv_transpose2d
.
padding
,
output_padding
=
conv_transpose2d
.
output_padding
,
dilation
=
conv_transpose2d
.
dilation
,
groups
=
conv_transpose2d
.
groups
,
bias
=
conv_transpose2d
.
bias
is
not
None
,
conv_mode
=
conv_transpose2d
.
conv_mode
,
compute_mode
=
conv_transpose2d
.
compute_mode
,
name
=
conv_transpose2d
.
name
,
)
new_conv_transpose2d
=
(
module
if
bn
is
None
or
not
conv_transpose2d
.
training
else
module
.
conv_transpose2d
)
weight
,
bias
=
conv_transpose2d
.
weight
,
conv_transpose2d
.
bias
if
not
conv_transpose2d
.
training
and
bn
is
not
None
:
weight
,
bias
=
fold_weight_bias
(
weight
,
bias
,
bn
.
weight
,
bn
.
bias
,
bn
.
running_mean
,
bn
.
running_var
,
bn
.
eps
,
transpose
=
False
,
)
new_conv_transpose2d
.
weight
=
Parameter
(
weight
)
if
bias
is
not
None
:
new_conv_transpose2d
.
bias
=
Parameter
(
bias
)
if
bn
is
not
None
and
conv_transpose2d
.
training
:
module
.
bn
=
deepcopy
(
bn
)
new_conv_transpose2d
.
training
=
conv_transpose2d
.
training
return
module
imperative/python/test/unit/module/test_qat.py
浏览文件 @
2694ff81
...
...
@@ -34,35 +34,49 @@ def test_qat_convbn2d():
in_channels
=
32
out_channels
=
64
kernel_size
=
3
class
TestNet
(
Module
):
def
__init__
(
self
,
groups
,
bias
):
super
().
__init__
()
self
.
quant
=
QuantStub
()
self
.
dequant
=
DequantStub
()
self
.
conv_bn
=
ConvBn2d
(
in_channels
,
out_channels
,
kernel_size
,
groups
=
groups
,
bias
=
bias
,
)
def
forward
(
self
,
inp
):
out
=
self
.
quant
(
inp
)
out
=
self
.
conv_bn
(
out
)
out
=
self
.
dequant
(
out
)
return
out
inputs
=
tensor
(
np
.
random
.
randn
(
4
,
in_channels
,
32
,
32
).
astype
(
np
.
float32
))
for
groups
,
bias
in
product
([
1
,
4
],
[
True
,
False
]):
module
=
ConvBn2d
(
in_channels
,
out_channels
,
kernel_size
,
groups
=
groups
,
bias
=
bias
)
M
.
init
.
normal_
(
module
.
bn
.
weight
)
M
.
init
.
normal_
(
module
.
bn
.
bias
)
module
.
train
()
qat_module
=
quantize_qat
(
module
,
inplace
=
False
)
disable_fake_quant
(
qat_module
)
inputs
=
tensor
(
np
.
random
.
randn
(
4
,
in_channels
,
32
,
32
).
astype
(
np
.
float32
))
normal_outputs
=
module
(
inputs
)
qat_outputs
=
qat_module
(
inputs
)
net
=
TestNet
(
groups
,
bias
)
net
.
train
()
qat_net
=
quantize_qat
(
net
,
inplace
=
False
)
disable_fake_quant
(
qat_net
)
normal_outputs
=
net
(
inputs
)
qat_outputs
=
qat_net
(
inputs
)
np
.
testing
.
assert_allclose
(
normal_outputs
.
numpy
(),
qat_outputs
.
numpy
(),
atol
=
5e-6
normal_outputs
.
numpy
(),
qat_outputs
.
numpy
(),
atol
=
1e-4
,
)
np
.
testing
.
assert_allclose
(
module
.
bn
.
running_mean
.
numpy
(),
qat_
module
.
bn
.
running_mean
.
numpy
(),
net
.
conv_bn
.
bn
.
running_mean
.
numpy
(),
qat_
net
.
conv_bn
.
bn
.
running_mean
.
numpy
(),
atol
=
5e-8
,
)
np
.
testing
.
assert_allclose
(
module
.
bn
.
running_var
.
numpy
(),
qat_module
.
bn
.
running_var
.
numpy
(),
atol
=
5e-7
,
net
.
conv_bn
.
bn
.
running_var
.
numpy
(),
qat_net
.
conv_bn
.
bn
.
running_var
.
numpy
(),
atol
=
5e-7
,
)
module
.
eval
()
normal_outputs
=
module
(
inputs
)
qat_
module
.
eval
()
qat_outputs
=
qat_
module
(
inputs
)
net
.
eval
()
normal_outputs
=
net
(
inputs
)
qat_
net
.
eval
()
qat_outputs
=
qat_
net
(
inputs
)
np
.
testing
.
assert_allclose
(
normal_outputs
.
numpy
(),
qat_outputs
.
numpy
(),
atol
=
5e-6
normal_outputs
.
numpy
(),
qat_outputs
.
numpy
(),
atol
=
1e-4
,
)
...
...
@@ -70,40 +84,44 @@ def test_qat_convtransposebn2d():
in_channels
=
32
out_channels
=
64
kernel_size
=
3
class
TestNet
(
Module
):
def
__init__
(
self
,
groups
,
bias
):
super
().
__init__
()
self
.
quant
=
QuantStub
()
self
.
dequant
=
DequantStub
()
self
.
conv_transpose_bn
=
ConvTransposeBn2d
(
in_channels
,
out_channels
,
kernel_size
,
groups
=
groups
,
bias
=
bias
,
)
def
forward
(
self
,
inp
):
out
=
self
.
quant
(
inp
)
out
=
self
.
conv_transpose_bn
(
out
)
out
=
self
.
dequant
(
out
)
return
out
for
groups
,
bias
in
product
([
1
,
4
],
[
True
,
False
]):
module
=
ConvTransposeBn2d
(
in_channels
=
in_channels
,
out_channels
=
out_channels
,
kernel_size
=
kernel_size
,
output_padding
=
0
,
groups
=
groups
,
bias
=
bias
,
)
M
.
init
.
normal_
(
module
.
bn
.
weight
)
M
.
init
.
normal_
(
module
.
bn
.
bias
)
module
.
train
()
qat_module
=
quantize_qat
(
module
,
inplace
=
False
)
disable_fake_quant
(
qat_module
)
net
=
TestNet
(
groups
,
bias
)
net
.
train
()
qat_net
=
quantize_qat
(
net
,
inplace
=
False
)
disable_fake_quant
(
qat_net
)
inputs
=
tensor
(
np
.
random
.
randn
(
4
,
in_channels
,
32
,
32
).
astype
(
np
.
float32
))
normal_outputs
=
module
(
inputs
)
qat_outputs
=
qat_module
(
inputs
)
np
.
testing
.
assert_allclose
(
normal_outputs
.
numpy
(),
qat_outputs
.
numpy
(),
atol
=
5e-6
)
normal_outputs
=
net
(
inputs
)
qat_outputs
=
qat_net
(
inputs
)
np
.
testing
.
assert_allclose
(
module
.
bn
.
running_mean
.
numpy
(),
qat_module
.
bn
.
running_mean
.
numpy
(),
atol
=
5e-8
,
normal_outputs
.
numpy
(),
qat_outputs
.
numpy
(),
atol
=
1e-5
,
)
np
.
testing
.
assert_allclose
(
module
.
bn
.
running_var
.
numpy
(),
qat_module
.
bn
.
running_var
.
numpy
(),
atol
=
5e-7
,
net
.
conv_transpose_bn
.
bn
.
running_var
.
numpy
(),
qat_net
.
conv_transpose_bn
.
bn
.
running_var
.
numpy
(),
atol
=
5e-7
,
)
module
.
eval
()
normal_outputs
=
module
(
inputs
)
qat_
module
.
eval
()
qat_outputs
=
qat_
module
(
inputs
)
net
.
eval
()
normal_outputs
=
net
(
inputs
)
qat_
net
.
eval
()
qat_outputs
=
qat_
net
(
inputs
)
np
.
testing
.
assert_allclose
(
normal_outputs
.
numpy
(),
qat_outputs
.
numpy
(),
atol
=
5e-6
normal_outputs
.
numpy
(),
qat_outputs
.
numpy
(),
atol
=
1e-5
,
)
...
...
imperative/python/test/unit/quantization/test_quantize.py
浏览文件 @
2694ff81
...
...
@@ -3,6 +3,15 @@ import pytest
from
megengine
import
Parameter
,
Tensor
from
megengine
import
module
as
Float
from
megengine.functional
import
ones
,
zeros
from
megengine.module
import
(
BatchNorm2d
,
Conv2d
,
ConvBn2d
,
ConvTranspose2d
,
ConvTransposeBn2d
,
ReLU
,
)
from
megengine.module
import
qat
as
QAT
from
megengine.module
import
quantized
as
Q
from
megengine.quantization
import
(
...
...
@@ -24,6 +33,7 @@ from megengine.quantization.quantize import (
quantize_qat
,
reset_qconfig
,
)
from
megengine.utils.bn_fusion
import
fuse_conv_bn_relu_module
class
FloatNet
(
Float
.
Module
):
...
...
@@ -291,3 +301,85 @@ def test_convert_with_custom_mapping():
net
=
Net
()
qat_net
=
quantize_qat
(
net
,
inplace
=
False
,
mapping
=
{
FloatExample
:
QATExample
})
assert
isinstance
(
qat_net
.
example
,
QATExample
)
def
test_ConvBn2d_fold_weight_bias
():
in_channels
=
32
out_channels
=
64
kernel_size
=
3
conv
=
Conv2d
(
in_channels
,
out_channels
,
kernel_size
)
bn
=
BatchNorm2d
(
out_channels
)
relu
=
ReLU
()
fused_conv
=
fuse_conv_bn_relu_module
(
conv
,
bn
,
relu
)
bn
.
eval
()
fused_conv
.
eval
()
inputs
=
Tensor
(
np
.
random
.
randn
(
4
,
in_channels
,
32
,
32
).
astype
(
np
.
float32
))
expected_result
=
relu
(
bn
(
conv
(
inputs
)))
actual_result
=
fused_conv
(
inputs
)
np
.
testing
.
assert_allclose
(
expected_result
.
numpy
(),
actual_result
.
numpy
(),
atol
=
1e-4
)
conv
.
eval
()
bn
.
eval
()
relu
.
eval
()
fused_conv
=
fuse_conv_bn_relu_module
(
conv
,
bn
,
relu
)
fused_conv
.
eval
()
expected_result
=
relu
(
conv
(
inputs
))
actual_result
=
fused_conv
(
inputs
)
np
.
testing
.
assert_allclose
(
expected_result
.
numpy
(),
actual_result
.
numpy
(),
atol
=
1e-4
)
conv
.
train
()
bn
.
train
()
fused_conv
=
fuse_conv_bn_relu_module
(
conv
,
bn
,
None
)
fused_conv
.
train
()
expected_result
=
bn
(
conv
(
inputs
))
actual_result
=
fused_conv
(
inputs
)
np
.
testing
.
assert_allclose
(
expected_result
.
numpy
(),
actual_result
.
numpy
(),
atol
=
1e-4
)
def
test_ConvTransposeBn2d_fold_weight_bias
():
in_channels
=
32
out_channels
=
64
kernel_size
=
3
conv
=
ConvTranspose2d
(
in_channels
,
out_channels
,
kernel_size
)
bn
=
BatchNorm2d
(
out_channels
)
relu
=
ReLU
()
fused_conv
=
fuse_conv_bn_relu_module
(
conv
,
bn
,
relu
)
bn
.
eval
()
fused_conv
.
eval
()
inputs
=
Tensor
(
np
.
random
.
randn
(
4
,
in_channels
,
32
,
32
).
astype
(
np
.
float32
))
expected_result
=
relu
(
bn
(
conv
(
inputs
)))
actual_result
=
fused_conv
(
inputs
)
np
.
testing
.
assert_allclose
(
expected_result
.
numpy
(),
actual_result
.
numpy
(),
atol
=
1e-4
)
conv
.
eval
()
bn
.
eval
()
relu
.
eval
()
fused_conv
=
fuse_conv_bn_relu_module
(
conv
,
bn
,
relu
)
fused_conv
.
eval
()
expected_result
=
relu
(
conv
(
inputs
))
actual_result
=
fused_conv
(
inputs
)
np
.
testing
.
assert_allclose
(
expected_result
.
numpy
(),
actual_result
.
numpy
(),
atol
=
1e-4
)
conv
.
train
()
bn
.
train
()
fused_conv
=
fuse_conv_bn_relu_module
(
conv
,
bn
,
None
)
fused_conv
.
train
()
expected_result
=
bn
(
conv
(
inputs
))
actual_result
=
fused_conv
(
inputs
)
np
.
testing
.
assert_allclose
(
expected_result
.
numpy
(),
actual_result
.
numpy
(),
atol
=
1e-4
)
编辑
预览
Markdown
is supported
0%
请重试
或
添加新附件
.
添加附件
取消
You are about to add
0
people
to the discussion. Proceed with caution.
先完成此消息的编辑!
取消
想要评论请
注册
或
登录