def test_get_img_shape_on_2d_image():
n = 5
channels = 4
dim1 = 1
dim2 = 2
K.set_image_data_format('channels_first')
assert (n, channels, dim1, dim2) == utils.get_img_shape(
K.ones(shape=(n, channels, dim1, dim2)))
K.set_image_data_format('channels_last')
assert (n, channels, dim1, dim2) == utils.get_img_shape(
K.ones(shape=(n, dim1, dim2, channels)))
def testConditionalMaskUpdate(self):
weight = K.variable(np.linspace(1.0, 100.0, 100), name="weights")
mask = K.ones(weight.get_shape())
threshold = K.zeros([])
def linear_sparsity(step):
sparsity_val = ops.convert_to_tensor(
[0.0, 0.1, 0.1, 0.3, 0.3, 0.5, 0.5, 0.5, 0.5, 0.5])
return ops.convert_to_tensor(True), sparsity_val[step]
# Set up pruning
p = pruning_impl.Pruning(pruning_vars=[(weight, mask, threshold)],
training_step_fn=self.training_step_fn,
pruning_schedule=linear_sparsity,
block_size=self.block_size,
block_pooling_type=self.block_pooling_type)
non_zero_count = []
for _ in range(10):
if context.executing_eagerly():
p.conditional_mask_update()
p.weight_mask_op()
state_ops.assign_add(self.global_step, 1)
else:
K.get_session().run(p.conditional_mask_update())
K.get_session().run(p.weight_mask_op())
K.get_session().run(state_ops.assign_add(self.global_step, 1))
non_zero_count.append(np.count_nonzero(K.get_value(weight)))
# Weights pruned at steps 1,3,5
expected_non_zero_count = [100, 90, 90, 70, 70, 50, 50, 50, 50, 50]
self.assertAllEqual(expected_non_zero_count, non_zero_count)
def call(self, inputs):
batch_size = K.shape(inputs)[0]
num_rows = K.int_shape(inputs)[1]
num_cols = K.int_shape(inputs)[2]
num_channels = K.int_shape(inputs)[3]
n = num_rows * num_cols
X = K.reshape(inputs, (batch_size, num_channels, n))
factor = K.cast(1 / n, K.floatx())
I_hat = factor * (K.eye(n) - factor * K.ones((n, n)))
I_hat = K.tile(
K.expand_dims(I_hat, axis=0),
(batch_size, 1, 1)) # One identity matrix per sample in batch
Sigma = K.batch_dot(K.batch_dot(X, I_hat),
K.permute_dimensions(X, (0, 2, 1)))
# Pre-normalization
trace = K.sum(K.sum(K.eye(num_channels) * Sigma, axis=1,
keepdims=True),
axis=2,
keepdims=True)
A = Sigma / trace
# Newton-Schulz Iteration
Y = A
Z = K.eye(num_channels)
Z = K.tile(K.expand_dims(Z, axis=0), (batch_size, 1, 1))
I3 = 3 * K.eye(num_channels)
I3 = K.tile(K.expand_dims(I3, axis=0), (batch_size, 1, 1))
for i in range(self.num_iter):
Y = 0.5 * K.batch_dot(Y, I3 - K.batch_dot(Z, Y))
Z = 0.5 * K.batch_dot(I3 - K.batch_dot(Z, Y), Z)
# Post-compensation
C = K.sqrt(trace) * Y
# Extract upper triangular matrix as vector
ones = K.ones((num_channels, num_channels))
mask = tf.matrix_band_part(ones, 0,
-1) # Upper triangular matrix of 0s and 1s
mask = K.cast(mask, 'bool') # Convert integer mask to boolean mask
triuvec = tf.boolean_mask(
C, mask, axis=1) # Apply mask to 2nd and 3rd dimension
triuvec.set_shape((None, num_channels * (num_channels + 1) /
2)) # Set correct shape manually
return triuvec
def call(self, inputs):
_, kernel_b = xnorize(self.kernel, self.H)
_, inputs_b = xnorize(inputs)
outputs = K.conv2d(inputs_b,
kernel_b,
strides=self.strides,
padding=self.padding,
data_format=self.data_format,
dilation_rate=self.dilation_rate)
# calculate Wa and xa
# kernel_a
mask = K.reshape(
self.kernel,
(-1,
self.filters)) # self.nb_row * self.nb_col * channels, filters
kernel_a = K.stop_gradient(K.mean(K.abs(mask), axis=0)) # filters
# inputs_a
if self.data_format == 'channels_first':
channel_axis = 1
else:
channel_axis = -1
mask = K.mean(K.abs(inputs), axis=channel_axis, keepdims=True)
ones = K.ones(self.kernel_size + (1, 1))
inputs_a = K.conv2d(mask,
ones,
strides=self.strides,
padding=self.padding,
data_format=self.data_format,
dilation_rate=self.dilation_rate
) # nb_sample, 1, new_nb_row, new_nb_col
if self.data_format == 'channels_first':
outputs = outputs * K.stop_gradient(inputs_a) * K.expand_dims(
K.expand_dims(K.expand_dims(kernel_a, 0), -1), -1)
else:
outputs = outputs * K.stop_gradient(inputs_a) * K.expand_dims(
K.expand_dims(K.expand_dims(kernel_a, 0), 0), 0)
if self.use_bias:
outputs = K.bias_add(outputs,
self.bias,
data_format=self.data_format)
if self.activation is not None:
return self.activation(outputs)
return outputs
def build(self, input_shape):
input_shape = tensor_shape.TensorShape(input_shape)
if len(input_shape) != 5:
raise ValueError('Inputs should have rank 5, '
'received input shape: %s' % input_shape)
if self.data_format == 'channels_first':
channel_axis = 1
else:
channel_axis = -1
if input_shape.dims[channel_axis].value is None:
raise ValueError('The channel dimension of the inputs '
'should be defined, found None: %s' % input_shape)
input_dim = int(input_shape[channel_axis])
self.input_spec = InputSpec(ndim=5, axes={channel_axis: input_dim})
if self.data_format == 'channels_first':
depth = int(input_shape[2])
else:
depth = int(input_shape[1])
kernel_shape = (depth, self.filter_size, self.filter_size, input_dim,
1)
# self.kernel = self.add_weight(
# 'kernel',
# shape=kernel_shape,
# initializer=self.kernel_initializer,
# regularizer=self.kernel_regularizer,
# constraint=self.kernel_constraint,
# trainable=False,
# dtype=self.dtype)
W = K.ones(kernel_shape, dtype=K.floatx())
W = W / K.cast(K.prod(K.int_shape(W)), dtype=K.floatx())
self.kernel = W
# self.set_weights([W])
if self.use_bias:
self.bias = self.add_weight('bias',
shape=(depth, self.filter_size,
self.filter_size),
initializer=self.bias_initializer,
regularizer=self.bias_regularizer,
constraint=self.bias_constraint,
trainable=False,
dtype=self.dtype)
else:
self.bias = None
self.built = True
def _blockMasking(self, block_size, block_pooling_type, weight,
expected_mask):
mask = K.ones(weight.get_shape())
threshold = K.zeros([])
# Set up pruning
p = pruning_impl.Pruning(pruning_vars=[(weight, mask, threshold)],
training_step_fn=self.training_step_fn,
pruning_schedule=self.constant_sparsity,
block_size=block_size,
block_pooling_type=block_pooling_type)
_, new_mask = p._maybe_update_block_mask(weight)
# Check if the mask is the same size as the weights
self.assertAllEqual(new_mask.get_shape(), weight.get_shape())
mask_after_pruning = K.get_value(new_mask)
self.assertAllEqual(mask_after_pruning, expected_mask)
def testUpdateSingleMask(self):
weight = K.variable(np.linspace(1.0, 100.0, 100), name="weights")
mask = K.ones(weight.get_shape())
threshold = K.zeros([])
p = pruning_impl.Pruning(pruning_vars=[(weight, mask, threshold)],
training_step_fn=self.training_step_fn,
pruning_schedule=self.constant_sparsity,
block_size=self.block_size,
block_pooling_type=self.block_pooling_type)
mask_before_pruning = K.get_value(mask)
self.assertAllEqual(np.count_nonzero(mask_before_pruning), 100)
if context.executing_eagerly():
p.conditional_mask_update()
else:
K.get_session().run(p.conditional_mask_update())
mask_after_pruning = K.get_value(mask)
self.assertAllEqual(np.count_nonzero(mask_after_pruning), 50)
def __call__(self, shape, dtype=None, partition_info=None):
# set bias to -log((1 - p)/p) for foreground
bias = -K.log((1 - self.probability) / self.probability)
result = K.get_value(K.ones(shape, dtype=dtype)) * bias
return result
def constant(input_batch, size):
batch_size = K.shape(input_batch)[0]
return K.tile(K.ones((1, size)), (batch_size, 1))
