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

Code before:
import tensorflow as tf


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 cluster_last_three_conv2d_layers(model):
  last_three_conv2d_layers =  [
      layer for layer in model.layers
      if isinstance(layer, tf.keras.layers.Conv2D) and
      not isinstance(layer, tf.keras.layers.DepthwiseConv2D)
    ]
  last_three_conv2d_layers = last_three_conv2d_layers[-3:]

  cluster_weights = tfmot.clustering.keras.cluster_weights
  CentroidInitialization = tfmot.clustering.keras.CentroidInitialization
  clustering_params1 = {
      'number_of_clusters': 256,
      'cluster_centroids_init': CentroidInitialization.LINEAR
  }
  clustering_params2 = {
      'number_of_clusters': 32,
      'cluster_centroids_init': CentroidInitialization.LINEAR
  }

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

    if layer == last_three_conv2d_layers[0] or layer == last_three_conv2d_layers[1]:
        clustered = cluster_weights(layer, **clustering_params1)
        print("Clustered {} with {} clusters".format(layer.name, clustering_params1['number_of_clusters']))
    else:
        clustered = cluster_weights(layer, **clustering_params2)
        print("Clustered {} with {} clusters".format(layer.name, clustering_params2['number_of_clusters']))
    return clustered

  return tf.keras.models.clone_model(model, clone_function=cluster_fn)


def run(flags_obj):
  """Run ResNet ImageNet training and eval loop using native Keras APIs.

Review comment (suggestion):
In my other comment, I suggested just directly inlining clustering_params1 and clustering_params2 with the call to `cluster_weights`. Was there a reason not to do so?

Inlining removes a layer of indirection.

Code after:
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 \
      layer == last_three_conv2d_layers[1]:
      clustered = cluster_weights(layer, number_of_clusters=256, \
          cluster_centroids_init=CentroidInitialization.LINEAR)
      print("Clustered {} with 256 clusters".format(layer.name))
    else:
      clustered = cluster_weights(layer, number_of_clusters=32, \
          cluster_centroids_init=CentroidInitialization.LINEAR)
      print("Clustered {} with 32 clusters".format(layer.name))
    return clustered

  return tf.keras.models.clone_model(model, clone_function=cluster_fn)


def run(flags_obj):
  """Run ResNet ImageNet training and eval loop using native Keras APIs.

  Args:
    flags_obj: An object containing parsed flag values.

  Raises:
    ValueError: If fp16 is passed as it is not currently supported.
    NotImplementedError: If some features are not currently supported.