text: #!/usr/bin/env python
  # -*- coding: utf-8 -*-
  # File: imagenet-resnet-keras.py
  # Author: Yuxin Wu <ppwwyyxxc@gmail.com>
  
  import numpy as np
  import tensorflow as tf
  import argparse
  
  from tensorpack import InputDesc, SyncMultiGPUTrainerReplicated
  from tensorpack.dataflow import FakeData, MapDataComponent
  from tensorpack.utils import logger
  from tensorpack.utils.gpu import get_nr_gpu
  from tensorpack.contrib.keras import KerasModel
  from tensorpack.callbacks import *
  from tensorflow.python.keras.layers import *
  
  from imagenet_utils import get_imagenet_dataflow, fbresnet_augmentor, ImageNetModel
  
  
  TOTAL_BATCH_SIZE = 512
  BASE_LR = 0.1 * (TOTAL_BATCH_SIZE // 256)
  
  
  def bn(x, name, zero_init=False):
      return BatchNormalization(
          axis=1, name=name, fused=True,
          momentum=0.9, epsilon=1e-5,
          gamma_initializer='zeros' if zero_init else 'ones')(x)
  
  
  def conv(x, filters, kernel, strides=1, name=None):
      return Conv2D(filters, kernel, name=name,
                    strides=strides, use_bias=False, padding='same',
                    kernel_initializer=tf.keras.initializers.VarianceScaling(
                        scale=2.0, mode='fan_out', distribution='normal'),
                    kernel_regularizer=tf.keras.regularizers.l2(5e-5))(x)
  
  
  def identity_block(input_tensor, kernel_size, filters, stage, block):
      filters1, filters2, filters3 = filters
      conv_name_base = 'res' + str(stage) + block + '_branch'
      bn_name_base = 'bn' + str(stage) + block + '_branch'
  
      x = conv(input_tensor, filters1, 1, name=conv_name_base + '2a')
      x = bn(x, name=bn_name_base + '2a')
      x = Activation('relu')(x)
  
      x = conv(x, filters2, kernel_size, name=conv_name_base + '2b')
      x = bn(x, name=bn_name_base + '2b')
      x = Activation('relu')(x)
  
      x = conv(x, filters3, (1, 1), name=conv_name_base + '2c')
      x = bn(x, name=bn_name_base + '2c', zero_init=True)
  
      x = tf.keras.layers.add([x, input_tensor])
      x = Activation('relu')(x)
      return x
  
  
  def conv_block(input_tensor, kernel_size, filters, stage, block, strides=(2, 2)):
      filters1, filters2, filters3 = filters
      conv_name_base = 'res' + str(stage) + block + '_branch'
      bn_name_base = 'bn' + str(stage) + block + '_branch'
  
      x = conv(input_tensor, filters1, (1, 1), name=conv_name_base + '2a')
      x = bn(x, name=bn_name_base + '2a')
      x = Activation('relu')(x)
  
      x = conv(x, filters2, kernel_size, strides=strides, name=conv_name_base + '2b')
      x = bn(x, name=bn_name_base + '2b')
      x = Activation('relu')(x)
  
      x = conv(x, filters3, (1, 1), name=conv_name_base + '2c')
      x = bn(x, name=bn_name_base + '2c', zero_init=True)
  
      shortcut = conv(
          input_tensor,
          filters3, (1, 1), strides=strides,
          name=conv_name_base + '1')
      shortcut = bn(shortcut, name=bn_name_base + '1')
  
      x = tf.keras.layers.add([x, shortcut])
      x = Activation('relu')(x)
      return x
  
  
  def resnet50(inputs):
      input = tf.layers.Input(tensor=inputs[0])
  
      def image_preprocess(image):
          image = ImageNetModel.image_preprocess(image)
          image = tf.transpose(image, [0, 3, 1, 2])
          return image
  
      x = Lambda(image_preprocess)(input)
  
      x = conv(x, 64, (7, 7), strides=(2, 2), name='conv0')
      x = bn(x, name='bn_conv1')
      x = Activation('relu')(x)
      x = MaxPooling2D((3, 3), strides=(2, 2), padding='same')(x)
  
      x = conv_block(x, 3, [64, 64, 256], stage=2, block='a', strides=(1, 1))
      x = identity_block(x, 3, [64, 64, 256], stage=2, block='b')
      x = identity_block(x, 3, [64, 64, 256], stage=2, block='c')
  
      x = conv_block(x, 3, [128, 128, 512], stage=3, block='a')
      x = identity_block(x, 3, [128, 128, 512], stage=3, block='b')
      x = identity_block(x, 3, [128, 128, 512], stage=3, block='c')
      x = identity_block(x, 3, [128, 128, 512], stage=3, block='d')
  
      x = conv_block(x, 3, [256, 256, 1024], stage=4, block='a')
      x = identity_block(x, 3, [256, 256, 1024], stage=4, block='b')
      x = identity_block(x, 3, [256, 256, 1024], stage=4, block='c')
      x = identity_block(x, 3, [256, 256, 1024], stage=4, block='d')
      x = identity_block(x, 3, [256, 256, 1024], stage=4, block='e')
      x = identity_block(x, 3, [256, 256, 1024], stage=4, block='f')
  
      x = conv_block(x, 3, [512, 512, 2048], stage=5, block='a')
      x = identity_block(x, 3, [512, 512, 2048], stage=5, block='b')
      x = identity_block(x, 3, [512, 512, 2048], stage=5, block='c')
  
      x = GlobalAveragePooling2D(name='avg_pool')(x)
      x = Flatten()(x)
      x = Dense(1000, activation='softmax', name='fc1000',
                kernel_initializer=tf.keras.initializers.VarianceScaling(
                    scale=2.0, mode='fan_in'),
                kernel_regularizer=tf.keras.regularizers.l2(5e-5))(x)
  
      M = tf.keras.models.Model(input, x, name='resnet50')
      return M
  
  
  if __name__ == '__main__':
      parser = argparse.ArgumentParser()
      parser.add_argument('--data', help='ILSVRC dataset dir')
      parser.add_argument('--fake', help='use fakedata to test or benchmark this model', action='store_true')
      args = parser.parse_args()
      logger.set_logger_dir("train_log/imagenet-resnet-keras")
  
      tf.keras.backend.set_image_data_format('channels_first')
  
      nr_gpu = get_nr_gpu()
      if args.fake:
          df_train = FakeData([[64, 224, 224, 3], [64, 1000]], 5000, random=False, dtype='uint8')
          df_val = FakeData([[64, 224, 224, 3], [64, 1000]], 5000, random=False)
      else:
          batch_size = TOTAL_BATCH_SIZE // nr_gpu
          assert args.data is not None
          df_train = get_imagenet_dataflow(
              args.data, 'train', batch_size, fbresnet_augmentor(True))
          df_val = get_imagenet_dataflow(
              args.data, 'val', batch_size, fbresnet_augmentor(False))
  
          def one_hot(label):
              return np.eye(1000)[label]
  
          df_train = MapDataComponent(df_train, one_hot, 1)
          df_val = MapDataComponent(df_val, one_hot, 1)
  
      M = KerasModel(
          resnet50,
          inputs_desc=[InputDesc(tf.uint8, [None, 224, 224, 3], 'images')],
          targets_desc=[InputDesc(tf.float32, [None, 1000], 'labels')],
          input=df_train,
          trainer=SyncMultiGPUTrainerReplicated(nr_gpu))
  
      lr = tf.get_variable('learning_rate', initializer=0.1, trainable=False)
      tf.summary.scalar('lr', lr)
  
      M.compile(
          optimizer=tf.train.MomentumOptimizer(lr, 0.9, use_nesterov=True),
          loss='categorical_crossentropy',
          metrics='categorical_accuracy'
      )
  
      callbacks = [
          ModelSaver(),
          ScheduledHyperParamSetter(
              'learning_rate',
              [(0, 0.1), (3, BASE_LR)], interp='linear'),  # warmup
          ScheduledHyperParamSetter(
              'learning_rate',
              [(30, BASE_LR * 0.1), (60, BASE_LR * 1e-2), (85, BASE_LR * 1e-3)]),
          GPUUtilizationTracker()
      ]
      if not args.fake:
          callbacks.append(
              DataParallelInferenceRunner(
                  df_val, ScalarStats(['categorical_accuracy']), nr_gpu))
  
      M.fit(
          steps_per_epoch=100 if args.fake else 1281167 // TOTAL_BATCH_SIZE,
          max_epoch=100,
          callbacks=callbacks
      )
  