def _basic_block( name: str, x_tens: tf_compat.Tensor, training: Union[bool, tf_compat.Tensor], out_channels: int, stride: int, kernel_initializer, bias_initializer, beta_initializer, gamma_initializer, ) -> tf_compat.Tensor: with tf_compat.variable_scope(name, reuse=tf_compat.AUTO_REUSE): out = conv2d_block( "conv_bn_0", x_tens, training, out_channels, kernel_size=3, stride=stride, padding=1, kernel_initializer=kernel_initializer, bias_initializer=bias_initializer, beta_initializer=beta_initializer, gamma_initializer=gamma_initializer, ) out = conv2d_block( "conv_bn_1", out, training, out_channels, kernel_size=3, padding=1, act=None, kernel_initializer=kernel_initializer, bias_initializer=bias_initializer, beta_initializer=beta_initializer, gamma_initializer=gamma_initializer, ) if stride > 1 or int(x_tens.shape[3]) != out_channels: out = tf_compat.add( out, _identity_modifier( x_tens, training, out_channels, stride, kernel_initializer=kernel_initializer, bias_initializer=bias_initializer, beta_initializer=beta_initializer, gamma_initializer=gamma_initializer, ), ) else: out = tf_compat.add(out, x_tens) out = activation(out, act="relu", name="act_out") return out
def cent_crop(img: tf_compat.Tensor): with tf_compat.name_scope(name): orig_shape = tf_compat.shape(img) min_size = tf_compat.cond( tf_compat.greater_equal(orig_shape[0], orig_shape[1]), lambda: orig_shape[1], lambda: orig_shape[0], ) if padding > 0: orig_shape_list = img.get_shape().as_list() resize((orig_shape_list[0] + 2 * padding, orig_shape_list[1] + 2 * padding)) padding_height = tf_compat.add( tf_compat.cast( tf_compat.round( tf_compat.div( tf_compat.cast( tf_compat.subtract(orig_shape[0], min_size), tf_compat.float32, ), 2.0, )), tf_compat.int32, ), padding, ) padding_width = tf_compat.add( tf_compat.cast( tf_compat.round( tf_compat.div( tf_compat.cast( tf_compat.subtract(orig_shape[1], min_size), tf_compat.float32, ), 2.0, )), tf_compat.int32, ), padding, ) img = tf_compat.image.crop_to_bounding_box(img, padding_height, padding_width, min_size, min_size) return img
def _inverted_bottleneck_block( name: str, x_tens: tf_compat.Tensor, training: Union[bool, tf_compat.Tensor], out_channels: int, exp_channels: int, stride: int, kernel_initializer, bias_initializer, beta_initializer, gamma_initializer, ) -> tf_compat.Tensor: with tf_compat.variable_scope(name, reuse=tf_compat.AUTO_REUSE): out = conv2d_block( "expand", x_tens, training, exp_channels, kernel_size=1, act="relu6", kernel_initializer=kernel_initializer, bias_initializer=bias_initializer, beta_initializer=beta_initializer, gamma_initializer=gamma_initializer, ) out = depthwise_conv2d_block( "spatial", out, training, exp_channels, kernel_size=3, stride=stride, act="relu6", kernel_initializer=kernel_initializer, bias_initializer=bias_initializer, beta_initializer=beta_initializer, gamma_initializer=gamma_initializer, ) out = conv2d_block( "compress", out, training, out_channels, kernel_size=1, act=None, kernel_initializer=kernel_initializer, bias_initializer=bias_initializer, beta_initializer=beta_initializer, gamma_initializer=gamma_initializer, ) if stride == 1 and int(x_tens.shape[3]) == out_channels: out = tf_compat.add(out, x_tens) return out
def _fc(name, x_tens, in_chan, out_chan, add_relu=True): with tf_compat.name_scope(name): weight = tf_compat.Variable( tf_compat.random_normal([in_chan, out_chan]), name="weight" ) bias = tf_compat.Variable(tf_compat.random_normal([out_chan]), "bias") x_tens = tf_compat.matmul(x_tens, weight, name="matmul") x_tens = tf_compat.add(x_tens, bias, name="add") if add_relu: x_tens = tf_compat.nn.relu(x_tens, name="relu") return x_tens
def test_create_op_pruning_fc(sparsity_val): group = "test-group" with tf_compat.Graph().as_default() as graph: inp = tf_compat.placeholder(tf_compat.float32, [None, 64]) with tf_compat.name_scope("fc"): weights = tf_compat.Variable(tf_compat.random_normal([64, 64]), name="weights") bias = tf_compat.Variable(tf_compat.random_normal([64]), name="bias") matmul = tf_compat.matmul(inp, weights, name="matmul") add = tf_compat.add(matmul, bias, name="bias_add") relu = tf_compat.nn.relu(add, name="relu") sparsity = tf_compat.placeholder(dtype=tf_compat.float32, name="sparsity_placeholder") update_ready = tf_compat.placeholder(dtype=tf_compat.bool, name="update_ready") matmul_op = graph.get_operation_by_name("fc/matmul") matmul_op_input = get_op_input_var(matmul_op, VAR_INDEX_FROM_TRAINABLE) pruning_op_vars = create_op_pruning( matmul_op, matmul_op_input, sparsity, update_ready, True, None, group, UnstructuredPruningMaskCreator(), ) with tf_compat.Session() as sess: sess.run(tf_compat.global_variables_initializer()) sess.run( pruning_op_vars.update, feed_dict={ sparsity: sparsity_val, update_ready: False }, ) mask_sparsity = eval_tensor_sparsity(pruning_op_vars.mask) weight_sparsity = eval_tensor_sparsity(pruning_op_vars.op_input) assert mask_sparsity < 1e-3 assert mask_sparsity == weight_sparsity masked_sparsity = eval_tensor_sparsity(pruning_op_vars.masked) assert masked_sparsity < 1e-3 sess.run( pruning_op_vars.update, feed_dict={ sparsity: sparsity_val, update_ready: True }, ) mask_sparsity = eval_tensor_sparsity(pruning_op_vars.mask) assert abs(mask_sparsity - sparsity_val) < 1e-3 masked_sparsity = eval_tensor_sparsity(pruning_op_vars.masked) assert abs(masked_sparsity - sparsity_val) < 1e-3 res = sess.run(relu, feed_dict={inp: numpy.random.random((4, 64))}) assert res.sum() > 0.0
def test_create_op_pruning_conv(sparsity_val: float, mask_creator: PruningMaskCreator): group = "test-group" is_grouped_mask = isinstance(mask_creator, GroupedPruningMaskCreator) with tf_compat.Graph().as_default() as graph: inp = tf_compat.placeholder(tf_compat.float32, [None, 8, 8, 64]) with tf_compat.name_scope("conv"): weights = tf_compat.Variable(tf_compat.random_normal( [3, 3, 64, 64]), name="weights") bias = tf_compat.Variable(tf_compat.random_normal([64]), name="bias") conv = tf_compat.nn.conv2d(inp, weights, strides=[1, 1, 1, 1], padding="SAME", name="conv") add = tf_compat.add(conv, bias, name="bias_add") relu = tf_compat.nn.relu(add, name="relu") sparsity = tf_compat.placeholder(dtype=tf_compat.float32, name="sparsity_placeholder") update_ready = tf_compat.placeholder(dtype=tf_compat.bool, name="update_ready") conv_op = graph.get_operation_by_name("conv/conv") conv_op_input = get_op_input_var(conv_op, VAR_INDEX_FROM_TRAINABLE) pruning_op_vars = create_op_pruning( conv_op, conv_op_input, sparsity, update_ready, True, None, group, mask_creator=mask_creator, ) with tf_compat.Session() as sess: sess.run(tf_compat.global_variables_initializer()) sess.run( pruning_op_vars.update, feed_dict={ sparsity: sparsity_val, update_ready: False }, ) err_threshold = 1e-3 if not is_grouped_mask else 0.05 mask_sparsity = eval_tensor_sparsity(pruning_op_vars.mask) weight_sparsity = eval_tensor_sparsity(pruning_op_vars.op_input) assert mask_sparsity < err_threshold assert abs(mask_sparsity - weight_sparsity) <= 1e-4 masked_sparsity = eval_tensor_sparsity(pruning_op_vars.masked) assert masked_sparsity < err_threshold sess.run( pruning_op_vars.update, feed_dict={ sparsity: sparsity_val, update_ready: True }, ) mask_sparsity = eval_tensor_sparsity(pruning_op_vars.mask) assert abs(mask_sparsity - sparsity_val) < err_threshold masked_sparsity = eval_tensor_sparsity(pruning_op_vars.masked) assert abs(masked_sparsity - sparsity_val) < err_threshold res = sess.run(relu, feed_dict={inp: numpy.random.random((4, 8, 8, 64))}) assert res.sum() > 0.0 if is_grouped_mask: # Check that every value in the mask_creator grouping # is the same within the mask. Assumes grouping applies # an absolte mean to each grouping grouped_mask = mask_creator.group_tensor(pruning_op_vars.mask) mask_vals_are_grouped = tf_compat.reduce_all( tf_compat.logical_or( tf_compat.equal(grouped_mask, 0.0), tf_compat.equal(grouped_mask, 1.0), )) assert sess.run(mask_vals_are_grouped)