Repository: tensorflow/models
File: official/benchmark/models/resnet_imagenet_main.py (Python)

Code before:
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================
"""Runs a ResNet model on the ImageNet dataset."""

from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

import os

from absl import app
from absl import flags
from absl import logging
import tensorflow as tf

from tensorflow_model_optimization.python.core.clustering.keras import cluster
import tensorflow_model_optimization as tfmot
from official.modeling import performance
from official.utils.flags import core as flags_core
from official.utils.misc import distribution_utils
from official.utils.misc import keras_utils
from official.utils.misc import model_helpers
from official.vision.image_classification import test_utils
from official.vision.image_classification.resnet import common
from official.vision.image_classification.resnet import imagenet_preprocessing
from official.vision.image_classification.resnet import resnet_model


def selective_layers_to_cluster(model):
  last_3conv2d_layers_to_cluster = [
      layer.name
      for layer in model.layers
      if isinstance(layer, tf.keras.layers.Conv2D) and
      not isinstance(layer, tf.keras.layers.DepthwiseConv2D)
  ]
  last_3conv2d_layers_to_cluster = last_3conv2d_layers_to_cluster[-3:]
  return last_3conv2d_layers_to_cluster


def selective_clustering_clone_wrapper(clustering_params1, clustering_params2,
                                       model):


Review comment (suggestion):
We didn't do this previously, but should start doing it now and the long term integration:

It'd be best to import `tfmot` only when the flags indicate that tfmot will be used, so that if there were hypothetically breakages, the rest of the benchmarks could continue to run.
   

So only if flags_obj.clustering_method

Code after:
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================
"""Runs a ResNet model on the ImageNet dataset."""

from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

import os

# Import libraries
from absl import app
from absl import flags
from absl import logging
import tensorflow as tf

from official.modeling import performance
from official.utils.flags import core as flags_core
from official.utils.misc import distribution_utils
from official.utils.misc import keras_utils
from official.utils.misc import model_helpers
from official.vision.image_classification import test_utils
from official.vision.image_classification.resnet import common
from official.vision.image_classification.resnet import imagenet_preprocessing
from official.vision.image_classification.resnet import resnet_model


def cluster_last_three_conv2d_layers(model):
  import tensorflow_model_optimization as tfmot
  last_three_conv2d_layers =  [
      layer for layer in model.layers
      if isinstance(layer, tf.keras.layers.Conv2D)
    ][-3:]

  cluster_weights = tfmot.clustering.keras.cluster_weights
  CentroidInitialization = tfmot.clustering.keras.CentroidInitialization

  def cluster_fn(layer):
    if layer not in last_three_conv2d_layers:
      return layer

    if layer == last_three_conv2d_layers[0] or \