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

Code before:

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):

  def apply_clustering_to_conv2d_but_depthwise(layer):
    layers_list = selective_layers_to_cluster(model)
    if layer.name in layers_list:
      if layer.name != layers_list[-1]:
        print("Wrapped layer " + layer.name +
              " with " +
              str(clustering_params1["number_of_clusters"]) + " clusters.")
        return cluster.cluster_weights(layer, **clustering_params1)
      else:
        print("Wrapped layer " + layer.name +
              " with " +
              str(clustering_params2["number_of_clusters"]) + " clusters.")

Review comment (suggestion):
I think these four functions (`selective_layers_to_cluster` to `get_selectively_clustered_model`) have unnecessary layers of indirection.

Something like this should suffice:

```
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)
     ]
     
     cluster_centroids_init = '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.cluster_weights(layer, clusters=256, cluster_centroids_init=cluster_centroids_init)
        else:
             clustered = cluster.cluster_weights(layer, clusters=32, cluster_centroids_init=cluster_centroids_init)

        print("Clustered {} with {} clusters".format(layer.name, clustered.number_of_clusters)

     return tf.keras.models.clone_model(model, clone_function=cluster_fn)
     
```
1. There isn't a need to have different function arguments unless they are actually configurable by an end user (which isn't needed in this current version ... only the long-term integration). 
2. See https://google.github.io/styleguide/pyguide.html#310-strings.
3. I switched to `layer` instead of `layer.name`, which I don't think there are disadvantages, and it makes the code a bit cleaner. 




Code after:

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 \
      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)