def _get_bn_params(sess: tf.compat.v1.Session,
bn: tf.Operation) -> libpymo.BNParams():
"""
helper to populate BN params from given BN op, required for fold
:param sess: tf.compat.v1.Session type
:param bn: BatchNorm or a FusedBatch Norm op
:return: bn_params
"""
# make sure you define the session and graph scope before loading vars from graph.
with sess.graph.as_default():
# create BNParams type and populate
bn_params = libpymo.BNParams()
bn_params.beta = BNUtils.get_beta_as_numpy_data(sess, bn).reshape(-1)
bn_params.gamma = BNUtils.get_gamma_as_numpy_data(sess, bn).reshape(-1)
bn_params.runningMean = BNUtils.get_moving_mean_as_numpy_data(
sess, bn).reshape(-1)
bn_params.runningVar = BNUtils.get_moving_variance_as_numpy_data(
sess, bn).reshape(-1)
epsilon = BNUtils.get_epsilon(bn)
var = BNUtils.get_moving_variance_as_numpy_data(sess, bn).reshape(-1)
var_with_epsilon = var + epsilon
sigma = np.sqrt(var_with_epsilon)
# sigma = tf.sqrt(BNUtils.get_moving_variance_as_numpy_data(sess, bn).reshape(-1) + epsilon)
bn_params.runningVar = sigma # sess.run(sigma).reshape(-1)
return bn_params
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