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