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