unmask the conv channels when sparse convolution is enabled
@@ -395,7 +395,8 @@ class DenseFeatureStackBlock(TrainableLayer):
stack = [input_tensor]
channel_dim = len(input_tensor.shape) - 1
- input_mask = tf.ones([input_tensor.shape.as_list()[-1]]) > 0
+ n_channels = input_tensor.shape.as_list()[-1]
+ input_mask = tf.ones([n_channels]) > 0
# Stack all convolution outputs
for idx, conv in enumerate(block.conv_layers):
@@ -415,22 +416,21 @@ class DenseFeatureStackBlock(TrainableLayer):
stack.append(conv)
- # unmask the stack
- channels = tf.unstack(stack, axis=-1)
- # exclude the input channels
- channels = channels[input_channels:]
- input_mask = input_mask[input_channels:]
- # create a channel with zeros to be placed where channels were not calculated
- channels.insert(0, tf.zeros_like(channels[0]))
- # re stack
- channels = tf.stack(channels, axis=-1)
+ if self.use_bdo: # unmask the conv channels
+ conv_channels = tf.concat(stack[1:], axis=-1)
+
+ # insert a channel with zeros to be placed
+ # where channels were not calculated
+ zero_channel = tf.zeros(conv_channels.shape[:-1])
+ zero_channel = tf.expand_dims(zero_channel, axis=-1)
+ conv_channels = tf.concat([zero_channel, conv_channels], axis=-1)
# indices to keep
- int_mask = tf.cast(input_mask, tf.int32)
+ int_mask = tf.cast(input_mask[n_channels:], tf.int32)
indices = tf.cumsum(int_mask) * int_mask
# rearrange stack with zeros where channels were not calculated
- stack = tf.gather(channels, indices, axis=-1)
- stack = tf.unstack(stack, axis=-1)
+ conv_channels = tf.gather(conv_channels, indices, axis=-1)
+ stack = [input_tensor, conv_channels]
return stack
