caffe BatchNorm#
from pathlib import Path
from google.protobuf import text_format
import caffe_pb2 as pb2
temp_dir = Path(".temp")
temp_dir.mkdir(exist_ok=True)
text = """
layer {
name: "data"
type: "Input"
top: "data"
input_param {
shape {
dim: 1
dim: 3
dim: 10
dim: 10
}
}
}
layer {
bottom: "data"
top: "bn"
name: "bn"
type: "BatchNorm"
batch_norm_param {
use_global_stats: true
}
}
layer {
bottom: "bn"
top: "bn"
name: "scale"
type: "Scale"
scale_param {
bias_term: true
}
}
"""
predict_net = text_format.Merge(text, pb2.NetParameter())
from tvm.relax.testing import nn
from tvm.relax import op as _op
exp_tab = {} # 存储节点
dtype = "float32"
# 优先处理输入层
for pl in predict_net.layer:
name = pl.name
if pl.type == "Input":
shape = pl.input_param.shape
assert len(shape)==1 and len(pl.top)==1, "Input 类型仅仅支持单输入单输出"
shape = list(shape[0].dim)
exp_tab[name] = nn.Placeholder(shape, dtype, name)
for pl in predict_net.layer:
name = pl.name
if pl.type == "BatchNorm":
assert len(pl.bottom) == 1
inp = exp_tab[pl.bottom[0]]
n, c, h, w = [int(sh) for sh in inp.struct_info.shape]
break
# exp_tab[name] = _op.concat(inputs, axis=pl.concat_param.axis)
exp_tab
{'data': data}
from tvm_book.frontend.caffe import check_unsupported_ops
supported_op_names = [
"BatchNorm",
"Concat",
"Convolution",
"Crop",
"Deconvolution",
"Dropout",
"Eltwise",
"Embed",
"Flatten",
"InnerProduct",
"Input",
"LRN",
"Permute",
"Pooling",
"Power",
"PReLU",
"ReLU",
"Reshape",
"Scale",
"Sigmoid",
"Slice",
"Softmax",
"TanH",
"Upsample",
"Reduction",
]
check_unsupported_ops(predict_net.layer, supported_op_names)
blob_file = "/media/pc/data/board/arria10/lxw/tasks/tools/npuusertools/models/caffe/resnet50/ResNet-50-model.caffemodel"
len(pl.blobs)
0
from caffe_fuse import fuse_network, get_bn_params
from caffe_utils import unity_struct
proto_file = "/media/pc/data/board/arria10/lxw/tasks/tools/npuusertools/models/caffe/resnet50/ResNet-50-deploy.prototxt"
blob_file = "/media/pc/data/board/arria10/lxw/tasks/tools/npuusertools/models/caffe/resnet50/ResNet-50-model.caffemodel"
# 加载网络定义和参数
init_net = pb2.NetParameter()
predict_net = pb2.NetParameter()
with open(proto_file, 'r') as f:
text_format.Merge(f.read(), predict_net)
with open(blob_file, 'rb') as fp:
init_net.ParseFromString(fp.read())
predict_net = unity_struct(predict_net)
init_net, predict_net = fuse_network(init_net, predict_net)
with open(temp_dir/"test.prototxt", "w") as fp: # 保存网络结构
fp.write(text_format.MessageToString(predict_net))
with open(temp_dir/"test.caffemodel", "wb") as fp: # 保存网络权重
fp.write(init_net.SerializeToString())
bn_layer = init_net.layer[1]
use_layer_field = bool(init_net.layer) # 判断用layer还是layers字段
init_layers = init_net.layer if use_layer_field else init_net.layers
init_layer_dict = {il.name: il for il in init_layers}
bn_params = get_bn_params(init_layer_dict, bn_layer)
_op.nn.batch_norm?
Signature:
_op.nn.batch_norm(
data: tvm.ir.expr.RelaxExpr,
gamma: tvm.ir.expr.RelaxExpr,
beta: tvm.ir.expr.RelaxExpr,
moving_mean: tvm.ir.expr.RelaxExpr,
moving_var: tvm.ir.expr.RelaxExpr,
axis: int,
epsilon: float = 1e-05,
center: bool = True,
scale: bool = True,
momentum: float = 0.1,
training: bool = True,
) -> tvm.ir.expr.RelaxExpr
Docstring:
Batch normalization layer (Ioffe and Szegedy, 2014).
Normalizes the input at each batch, i.e. applies a transformation
that maintains the mean activation close to 0 and the activation
standard deviation close to 1.
.. math::
data\_mean[i] = mean(data[:,i,:,...]) \\
data\_var[i] = var(data[:,i,:,...])
Both *mean* and *var* returns a scalar by treating the input as a vector.
Then compute the normalized output, which has the same shape as input, as following:
.. math::
out[:,i,:,...] = \frac{data[:,i,:,...] - data\_mean[i]}{\sqrt{data\_var[i]+\epsilon}}
* gamma[i] + beta[i]
Assume the input has size *k* on axis 1, then both ``gamma`` and ``beta``
have shape *(k,)*.
Besides the inputs and the outputs, this operator accepts two auxiliary
states, ``moving_mean`` and ``moving_var``, which are *k*-length
vectors. They are global statistics for the whole dataset, which are updated by
.. code:: python
moving_mean = moving_mean * momentum + data_mean * (1 - momentum)
moving_var = moving_var * momentum + data_var * (1 - momentum)
The parameter ``axis`` specifies which axis of the input shape denotes
the 'channel' (separately normalized groups). The default is 1.
Specifying -1 sets the channel axis to be the last item in the input shape.
.. note::
This operator has two modes:
- Training mode.
- Use the mean and var computed from THIS batch to normalize.
- Update and then return the running mean and running var.
- Inference mode.
- Use the running_mean and running_var parameters to normalize.
- Do not update the running mean and running var. Just return the original value.
In the legalization stage, this operator will be legalized to the training mode by default.
You can use tvm.relax.transform.DecomposeOpsForInference to decompose the operator, so it
executes the inference mode computation. Similarly, use
tvm.relax.transform.DecomposeOpsForTraining to execute the training mode computation.
Parameters
----------
data : relax.Expr
The input data to the operator.
gamma : relax.Expr
The gamma scale factor.
beta : relax.Expr
The beta offset factor.
moving_mean : relax.Expr
Running mean of input.
moving_var : relax.Expr
Running variance of input.
axis : int
The axis along which the normalization is applied.
epsilon : float
Small float added to variance to avoid dividing by zero.
center : bool
Indicating if the beta offset will be added to the normalized tensor.
scale : bool
Indicating if the gamma scale will be multiplied.
momentum : float
The value used for the moving_mean and moving_var update.
training : bool
A boolean value to indicate whether training or in eval mode. By default.
relax batch_norm is training mode. To transform it to inference mode,
can use DecomposeOpsForInference.
Returns
-------
result : relax.Expr
The computed result.
File: /media/pc/data/lxw/ai/tvm/python/tvm/relax/op/nn/nn.py
Type: function
len(bn_layer.blobs)
3