federated_model.py

# Copyright 2021, Lanping-Tech.



"""ResNet v2 model for Keras using Batch or Group Normalization.

Related papers/blogs:
- http://arxiv.org/pdf/1603.05027v2.pdf

"""
from typing import Optional

import tensorflow as tf
import tensorflow_addons.layers.normalizations as tfa_norms

from tensorflow.keras import layers, models, losses, metrics
import tensorflow_federated as tff

import importlib

BATCH_NORM_DECAY = 0.997
BATCH_NORM_EPSILON = 1e-5
L2_WEIGHT_DECAY = 1e-4


def _norm_relu(input_tensor, norm='group'):
  """Helper function to make a Norm -> ReLU block."""
  if tf.keras.backend.image_data_format() == 'channels_last':
    channel_axis = 3
  else:
    channel_axis = 1

  if norm == 'group':
    x = tfa_norms.GroupNormalization(axis=channel_axis)(input_tensor)
  else:
    x = tf.keras.layers.BatchNormalization(
        axis=channel_axis,
        momentum=BATCH_NORM_DECAY,
        epsilon=BATCH_NORM_EPSILON)(input_tensor)
  return tf.keras.layers.Activation('relu')(x)


def _conv_norm_relu(input_tensor,
                    filters,
                    kernel_size,
                    strides=(1, 1),
                    norm='group',
                    seed: Optional[int] = None):
  """Helper function to make a Conv -> Norm -> ReLU block."""
  x = tf.keras.layers.Conv2D(
      filters,
      kernel_size,
      strides=strides,
      padding='same',
      use_bias=False,
      kernel_initializer=tf.keras.initializers.HeNormal(seed=seed),
      kernel_regularizer=tf.keras.regularizers.l2(L2_WEIGHT_DECAY))(
          input_tensor)
  return _norm_relu(x, norm=norm)


def _norm_relu_conv(input_tensor,
                    filters,
                    kernel_size,
                    strides=(1, 1),
                    norm='group',
                    seed: Optional[int] = None):
  """Helper function to make a Norm -> ReLU -> Conv block."""
  x = _norm_relu(input_tensor, norm=norm)
  x = tf.keras.layers.Conv2D(
      filters,
      kernel_size,
      strides=strides,
      padding='same',
      use_bias=False,
      kernel_initializer=tf.keras.initializers.HeNormal(seed=seed),
      kernel_regularizer=tf.keras.regularizers.l2(L2_WEIGHT_DECAY))(
          x)
  return x


def _shortcut(input_tensor, residual, norm='group', seed: Optional[int] = None):
  """Adds a shortcut between input and the residual."""
  input_shape = tf.keras.backend.int_shape(input_tensor)
  residual_shape = tf.keras.backend.int_shape(residual)

  if tf.keras.backend.image_data_format() == 'channels_last':
    row_axis = 1
    col_axis = 2
    channel_axis = 3
  else:
    channel_axis = 1
    row_axis = 2
    col_axis = 3

  stride_width = int(round(input_shape[row_axis] / residual_shape[row_axis]))
  stride_height = int(round(input_shape[col_axis] / residual_shape[col_axis]))
  equal_channels = input_shape[channel_axis] == residual_shape[channel_axis]

  shortcut = input_tensor
  # 1 X 1 conv if shape is different. Else identity.
  if stride_width > 1 or stride_height > 1 or not equal_channels:
    shortcut = tf.keras.layers.Conv2D(
        filters=residual_shape[channel_axis],
        kernel_size=(1, 1),
        strides=(stride_width, stride_height),
        padding='valid',
        use_bias=False,
        kernel_initializer=tf.keras.initializers.HeNormal(seed=seed),
        kernel_regularizer=tf.keras.regularizers.l2(L2_WEIGHT_DECAY))(
            shortcut)

    if norm == 'group':
      shortcut = tfa_norms.GroupNormalization(axis=channel_axis)(shortcut)
    else:
      shortcut = tf.keras.layers.BatchNormalization(
          axis=channel_axis,
          momentum=BATCH_NORM_DECAY,
          epsilon=BATCH_NORM_EPSILON)(shortcut)

  return tf.keras.layers.add([shortcut, residual])


def _basic_block(input_tensor,
                 filters,
                 strides=(1, 1),
                 avoid_norm=False,
                 norm='group',
                 seed: Optional[int] = None):
  """Basic convolutional block for use on resnets with <= 34 layers."""
  if avoid_norm:
    x = tf.keras.layers.Conv2D(
        filters=filters,
        kernel_size=(3, 3),
        strides=strides,
        padding='same',
        use_bias=False,
        kernel_initializer=tf.keras.initializers.HeNormal(seed=seed),
        kernel_regularizer=tf.keras.regularizers.l2(L2_WEIGHT_DECAY))(
            input_tensor)
  else:
    x = _norm_relu_conv(
        input_tensor,
        filters=filters,
        kernel_size=(3, 3),
        strides=strides,
        norm=norm,
        seed=seed)

  x = _norm_relu_conv(
      x,
      filters=filters,
      kernel_size=(3, 3),
      strides=strides,
      norm=norm,
      seed=seed)
  return _shortcut(input_tensor, x, norm=norm, seed=seed)


def _bottleneck_block(input_tensor,
                      filters,
                      strides=(1, 1),
                      avoid_norm=False,
                      norm='group',
                      seed: Optional[int] = None):
  """Bottleneck convolutional block for use on resnets with > 34 layers."""
  if avoid_norm:
    x = tf.keras.layers.Conv2D(
        filters=filters,
        kernel_size=(1, 1),
        strides=strides,
        padding='same',
        use_bias=False,
        kernel_initializer=tf.keras.initializers.HeNormal(seed=seed),
        kernel_regularizer=tf.keras.regularizers.l2(L2_WEIGHT_DECAY))(
            input_tensor)
  else:
    x = _norm_relu_conv(
        input_tensor,
        filters=filters,
        kernel_size=(1, 1),
        strides=strides,
        norm=norm,
        seed=seed)

  x = _norm_relu_conv(
      x,
      filters=filters,
      kernel_size=(3, 3),
      strides=strides,
      norm=norm,
      seed=seed)

  x = _norm_relu_conv(
      x,
      filters=filters * 4,
      kernel_size=(1, 1),
      strides=strides,
      norm=norm,
      seed=seed)
  return _shortcut(input_tensor, x, norm=norm, seed=seed)


def _residual_block(input_tensor,
                    block_function,
                    filters,
                    num_blocks,
                    strides=(1, 1),
                    is_first_layer=False,
                    norm='group',
                    seed: Optional[int] = None):
  """Builds a residual block with repeating bottleneck or basic blocks."""
  x = input_tensor
  for i in range(num_blocks):
    avoid_norm = is_first_layer and i == 0
    x = block_function(
        x,
        filters=filters,
        strides=strides,
        avoid_norm=avoid_norm,
        norm=norm,
        seed=seed)
  return x


def create_resnet(input_shape,
                  num_classes=10,
                  block='bottleneck',
                  repetitions=None,
                  initial_filters=64,
                  initial_strides=(2, 2),
                  initial_kernel_size=(7, 7),
                  initial_pooling='max',
                  norm='group',
                  seed: Optional[int] = None):
  """Instantiates a ResNet v2 model with Group Normalization.

  Instantiates the architecture from http://arxiv.org/pdf/1603.05027v2.pdf.
  The ResNet contains stages of residual blocks. Each residual block contains
  some number of...

  Args:
    input_shape: A tuple of length 3 describing the number of rows, columns, and
    channels of an input. Can be in channel-first or channel-last format.
    num_classes: Number of output classes.
    block: Whether to use a bottleneck or basic block within each stage.
    repetitions: A list of integers describing the number of blocks within each
      stage. If None, defaults to the resnet50 repetitions of [3, 4, 6, 3].
    initial_filters: The number of filters in the initial conv layer.
    initial_strides: The strides in the initial conv layer.
    initial_kernel_size: The kernel size for the initial conv layer.
    initial_pooling: The type of pooling after the initial conv layer.
    norm: Type of normalization to be used. Can be 'group' or 'batch'.
    seed: A random seed governing the model initialization and layer randomness.
      If not `None`, then the global random seed will be set before constructing
      the tensor initializer, in order to guarantee the same model is produced.

  Returns:
    A `tf.keras.Model`.

  Raises:
    Exception: Input shape should be a tuple of length 3.
  """
  if seed is not None:
    tf.random.set_seed(seed)

  if len(input_shape) != 3:
    raise Exception(
        'Input shape should be a tuple of length 3.')
  if repetitions is None:
    repetitions = [3, 4, 6, 3]

  if block == 'basic':
    block_fn = _basic_block
  elif block == 'bottleneck':
    block_fn = _bottleneck_block

  img_input = tf.keras.layers.Input(shape=input_shape)
  x = _conv_norm_relu(
      img_input,
      filters=initial_filters,
      kernel_size=initial_kernel_size,
      strides=initial_strides,
      norm=norm,
      seed=seed)

  if initial_pooling == 'max':
    x = tf.keras.layers.MaxPooling2D(
        pool_size=(3, 3), strides=initial_strides, padding='same')(x)

  filters = initial_filters

  for i, r in enumerate(repetitions):
    x = _residual_block(
        x,
        block_fn,
        filters=filters,
        num_blocks=r,
        is_first_layer=(i == 0),
        norm=norm,
        seed=seed)
    filters *= 2

  # Final activation in the residual blocks
  x = _norm_relu(x, norm=norm)

  # Classification block
  x = tf.keras.layers.GlobalAveragePooling2D()(x)

  x = tf.keras.layers.Dense(
      num_classes,
      activation='softmax',
      kernel_initializer=tf.keras.initializers.RandomNormal(
          stddev=0.01, seed=seed),
      kernel_regularizer=tf.keras.regularizers.l2(L2_WEIGHT_DECAY),
      bias_regularizer=tf.keras.regularizers.l2(L2_WEIGHT_DECAY))(
          x)

  model = tf.keras.models.Model(img_input, x)
  return model


def create_resnet18(input_shape,
                    num_classes,
                    norm='group',
                    seed: Optional[int] = None):
  """ResNet with 18 layers and basic residual blocks."""
  return create_resnet(
      input_shape,
      num_classes,
      'basic',
      repetitions=[2, 2, 2, 2],
      norm=norm,
      seed=seed)


def create_resnet34(input_shape,
                    num_classes,
                    norm='group',
                    seed: Optional[int] = None):
  """ResNet with 34 layers and basic residual blocks."""
  return create_resnet(
      input_shape,
      num_classes,
      'basic',
      repetitions=[3, 4, 6, 3],
      norm=norm,
      seed=seed)


def create_resnet50(input_shape,
                    num_classes,
                    norm='group',
                    seed: Optional[int] = None):
  """ResNet with 50 layers and bottleneck residual blocks."""
  return create_resnet(
      input_shape,
      num_classes,
      'bottleneck',
      repetitions=[3, 4, 6, 3],
      norm=norm,
      seed=seed)


def create_resnet101(input_shape,
                     num_classes,
                     norm='group',
                     seed: Optional[int] = None):
  """ResNet with 101 layers and bottleneck residual blocks."""
  return create_resnet(
      input_shape,
      num_classes,
      'bottleneck',
      repetitions=[3, 4, 23, 3],
      norm=norm,
      seed=seed)


def create_resnet152(input_shape,
                     num_classes,
                     norm='group',
                     seed: Optional[int] = None):
  """ResNet with 152 layers and bottleneck residual blocks."""
  return create_resnet(
      input_shape,
      num_classes,
      'bottleneck',
      repetitions=[3, 8, 36, 3],
      norm=norm,
      seed=seed)


def get_federated_model_from_keras(model_name, input_shape, input_spec, n_class):
	keras_model = getattr(importlib.import_module('federated_model'),'create_'+model_name)
	def model_fn():
		return tff.learning.from_keras_model(
			keras_model(input_shape, n_class),
			input_spec=input_spec,
			loss=losses.CategoricalCrossentropy(),
			metrics=[metrics.CategoricalAccuracy()])
	return model_fn

if __name__ =="__main__":
	getattr(importlib.import_module('federated_model'),'create_resnet18')((32,32,3),10).summary()