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)