MBConvBlockWithoutDepthwise stride implemented in 1x1 projection, wasting expansion arithmetic

[andravin]

MBConvBlockWithoutDepthwise implements stride in the 1x1 projection convolution. When stride=2, the projection discards 3/4ths of the activations produced by the expansion. It would be equivalent to implement stride on the 3x3 expansion convolution instead, and this would reduce the total block arithmetic almost by a factor of 4.

tpu/models/official/efficientnet/efficientnet_model.py

Lines 422 to 442 in 8462d08

	self._expand_conv = tf.layers.Conv2D(
	filters,
	kernel_size=[3, 3],
	strides=[1, 1],
	kernel_initializer=conv_kernel_initializer,
	padding='same',
	use_bias=False)
	self._bn0 = self._batch_norm(
	axis=self._channel_axis,
	momentum=self._batch_norm_momentum,
	epsilon=self._batch_norm_epsilon)
	# Output phase:
	filters = self._block_args.output_filters
	self._project_conv = tf.layers.Conv2D(
	filters,
	kernel_size=[1, 1],
	strides=self._block_args.strides,
	kernel_initializer=conv_kernel_initializer,
	padding='same',
	use_bias=False)

[andravin]

The reduction in OPs is substantial: moving the stride from the projection convolution to the expansion reduces the total number of operations in EfficientEdgeTPU-S by about 24%.

[andravin]

At this point, it should be apparent that Google has no intention of responding to this issue. I will nevertheless leave this comment here in case it helps anyone who should stumble upon it:

Google actually fixed this design flaw in the updated implementation of the MBConWithoutDepthwise block as it is used in MobileDets. In that work, they renamed the block "Fused convolution layer."

The implementation can be found here: https://github.com/tensorflow/models/blob/2986bcafb9eaa8fed4d78f17a04c4c5afc8f6691/research/object_detection/models/ssd_mobiledet_feature_extractor.py#L142-L147

Notice that the stride is now implemented in the expansion convolution:

    h = _conv(h, expanded_filters, kernel_size, strides=strides,
              activation_fn=activation_fn)
    if use_se:
      hidden_dim = _scale_filters(expanded_filters, 0.25)
      h = _squeeze_and_excite(h, hidden_dim, activation_fn=activation_fn)
    h = _conv(h, filters, 1, activation_fn=tf.identity)

[mingxingtan]

Hi @andravin Thanks for the great point. I have prepared a fix change internally, which should go out in the next release.

[allenwang28]

Fixed in 32572cb

[yuezha01]

I noticed exactly the same issue recently when I looked into the Efficientnet_EdgeTPU repo. Thank @andravin for pointing out the issue here.
Quick note here in case anyone uses the TFlite files (e.g. EfficientNet-EdgeTPU-S) published in the efficientnet_edgeTPU repo. Although the issue has been fixed in the new released code, the Tflite files still has stride of 2 in the point wise layer of fused block. As @andravin mentioned above, the operation counts measured for this published TFlite is unnecessarily higher (~30%).