def test_bn_fold(self):
# Generating random numbers from a normal distribution for the weights and biases of the current and prev layer
np.random.seed(1)
total = 2 * 3 * 2 * 2
beta = np.array(np.random.randn(2))
gamma = np.array(np.random.randn(2))
running_mean = np.array(np.random.randn(2))
running_var = np.array(np.random.randn(2))
bn_params = libpymo.BNParams()
bn_params.beta = beta
bn_params.gamma = gamma
bn_params.runningMean = running_mean
bn_params.runningVar = running_var
weight_tensor = libpymo.TensorParams()
weight = np.array(np.random.randn(total))
weight_sz = np.array([2, 3, 2, 2])
weight_tensor.data = weight
weight_tensor.shape = weight_sz
random_bias = np.array(np.random.rand(2))
bias_tensor = libpymo.TensorParams()
bias_tensor.data = random_bias
bias_tensor.shape = np.array([2])
w, b = bn_fold_prev_layer(weight.reshape(weight_sz), random_bias, beta,
gamma, running_mean, running_var)
bias = libpymo.fold(bn_params, weight_tensor, bias_tensor, True, True)
assert (np.allclose(w.flatten(), weight_tensor.data))
assert (np.allclose(b.flatten(), bias))
def test_bn_fold_to_next_linear_layer(self):
np.random.seed(1)
total = 4 * 2
weight = np.array(np.random.randn(total))
weight_sz = np.array([4, 2, 1, 1])
beta = np.array(np.random.randn(2))
gamma = np.array(np.random.randn(2))
running_mean = np.array(np.random.randn(2))
running_var = np.array(np.random.randn(2))
bn_params = libpymo.BNParams()
bn_params.beta = beta
bn_params.gamma = gamma
bn_params.runningMean = running_mean
bn_params.runningVar = running_var
layer_weight_params = libpymo.TensorParams()
layer_weight_params.data = weight
layer_weight_params.shape = weight_sz
random_bias = np.array(np.random.rand(4))
bias_tensor = libpymo.TensorParams()
bias_tensor.data = random_bias
bias_tensor.shape = np.array([4])
w, b = bn_fold_next_layer(weight.reshape(weight_sz), random_bias, beta,
gamma, running_mean, running_var)
bias = libpymo.fold(bn_params, layer_weight_params, bias_tensor, True,
False)
assert (np.allclose(w.flatten(), layer_weight_params.data))
assert (np.allclose(b.flatten(), bias))
def call_mo_batch_norm_fold(
conv_linear: Union[torch.nn.Linear, torch.nn.Conv2d,
torch.nn.ConvTranspose2d], bn: torch.nn.BatchNorm2d,
is_batch_norm_second: bool
) -> [torch.nn.Parameter, torch.nn.Parameter]:
"""
Calls Model optimization batch norm fold code and returns updated bias and weight
:param conv_linear: Conv or Linear layer. For Conv layers Conv2D and TransposedConv2D are supported currently
:param bn: Batch Norm layer
:param is_batch_norm_second: True if BatchNorm comes after Conv/Linear layer
:return: Updated bias and weight
"""
bn_params = libpymo.BNParams()
bn_params.gamma = bn.weight.detach().numpy().reshape(-1)
bn_params.beta = bn.bias.detach().numpy().reshape(-1)
bn_params.runningMean = bn.running_mean.detach().numpy().reshape(-1)
sigma = torch.sqrt(bn.running_var + bn.eps)
bn_params.runningVar = sigma.detach().numpy().reshape(-1)
weight_tensor = libpymo.TensorParams()
weight = conv_linear.weight
# Transpose weights to C, N, H, W from N, C, H, W since axis are flipped for transposed conv
# However depthwise conv layers are always N, 1, H, W whether transposed-conv or not, so no need to transpose
if isinstance(conv_linear,
torch.nn.ConvTranspose2d) and conv_linear.groups == 1:
weight = weight.permute(1, 0, 2, 3)
weight_tensor.data = weight.detach().numpy().reshape(-1)
weight_shape = np.array(weight.shape)
if len(conv_linear.weight.shape) == 2:
weight_shape = np.append(weight_shape, [1, 1])
weight_tensor.shape = weight_shape
bias_tensor = libpymo.TensorParams()
is_bias_valid = False
if conv_linear.bias is not None:
bias_tensor.data = conv_linear.bias.detach().numpy().reshape(-1)
bias_tensor.shape = np.array(conv_linear.bias.shape)
is_bias_valid = True
bias = libpymo.fold(bn_params, weight_tensor, bias_tensor, is_bias_valid,
is_batch_norm_second)
return bias, weight_tensor
def _fold_given_auto_selected_batch_norms(
sess: tf.compat.v1.Session,
layer_pairs: List[PairType]) -> tf.compat.v1.Session:
"""
Fold a given set of batch_norm layers into conv layers
:param sess: tf.compat.v1.Session
:param layer_pairs: pair of conv and bn layers
:return: new session with updated graph
"""
with sess.graph.as_default():
for pair in layer_pairs:
conv_linear, batchnorm, is_batch_norm_second = pair
assert conv_linear.type in [
'Conv2D', 'DepthwiseConv2dNative', 'MatMul'
]
# check flag
is_bias_valid = False
if not BiasUtils.is_bias_none(conv_linear):
is_bias_valid = True
bn_params = _get_bn_params(sess, batchnorm.op)
weight_tensor = _get_weight_tensor_transpose_reshape(
sess, conv_linear)
bias_tensor = _get_bias_tensor(sess, conv_linear)
bias = libpymo.fold(bn_params, weight_tensor, bias_tensor,
is_bias_valid, is_batch_norm_second)
# converting back to TF format [kh, kw, Nic, Noc] before updating weight tensor value
if conv_linear.type == 'DepthwiseConv2dNative':
# Depthwise conv layers in TF have outputs(Noc) set to 1.
# we send in format [Nic, Noc, kh, kw]
numpy_weight_reshaped = np.reshape(
weight_tensor.data, weight_tensor.shape).transpose(
(2, 3, 0, 1))
elif conv_linear.type == 'MatMul':
# o, i - convert to i , o
numpy_weight_reshaped = np.reshape(
weight_tensor.data,
[weight_tensor.shape[0], weight_tensor.shape[1]
]).transpose(1, 0)
else:
# conv2D case
# we sent in format [Noc, Nic, kh, kw]
numpy_weight_reshaped = np.reshape(
weight_tensor.data, weight_tensor.shape).transpose(
(2, 3, 1, 0))
WeightTensorUtils.update_tensor_for_op(sess, conv_linear,
numpy_weight_reshaped)
# remove bn op
BNUtils.skip_bn_op(sess, batchnorm.op, batchnorm.in_tensor,
batchnorm.out_tensor)
# update bias tensor, even in case there was no existing bias add op in given conv2D op.
bias_tensor_shape = [weight_tensor.shape[0]]
numpy_bias_reshaped = np.reshape(bias, bias_tensor_shape)
BiasUtils.update_bias_for_op(sess, conv_linear,
numpy_bias_reshaped)
# we edited the graph, so we should load and save for the metagraph associated with the session to be updated
after_bn_fold_sess = save_and_load_graph('./temp_bn_fold', sess)
return after_bn_fold_sess