All from scratch

[leondgarse]

Beforehand Data Prepare

---

Data source from Insightface

• Training Data in this project is MS1M-ArcFace downloaded from Insightface Dataset Zoo
• Evaluating data is LFW CFP-FP AgeDB-30 bin files included in MS1M-ArcFace dataset
• Any other data is also available just in the right format

Training dataset

• Extract data from mxnet record format to folders like

  .
  ├── 0
  │   ├── 100.jpg
  │   ├── 101.jpg
  │   └── 102.jpg
  ├── 1
  │   ├── 111.jpg
  │   ├── 112.jpg
  │   └── 113.jpg
  ├── 10
  │   ├── 707.jpg
  │   ├── 708.jpg
  │   └── 709.jpg

• Python code source folder is /datasets/faces_emore, extract to dist folder /datasets/faces_emore_112x112_folders, this may take some time.

  import os
  import numpy as np
  import mxnet as mx
  from tqdm import tqdm
  
  read_dir = '/datasets/faces_emore'
  save_dir = '/datasets/faces_emore_112x112_folders'
  idx_path = os.path.join(read_dir, 'train.idx')
  bin_path = os.path.join(read_dir, 'train.rec')
  
  imgrec = mx.recordio.MXIndexedRecordIO(idx_path, bin_path, 'r')
  rec_header, _ = mx.recordio.unpack(imgrec.read_idx(0))
  
  for ii in tqdm(range(1, int(rec_header.label[0]))):
      img_info = imgrec.read_idx(ii)
      header, img = mx.recordio.unpack(img_info)
      img_idx = str(int(np.sum(header.label)))
      img_save_dir = os.path.join(save_dir, img_idx)
      if not os.path.exists(img_save_dir):
          os.makedirs(img_save_dir)
      # print(os.path.join(img_save_dir, str(ii) + '.jpg'))
      with open(os.path.join(img_save_dir, str(ii) + '.jpg'), 'wb') as ff:
          ff.write(img)

Evaluating bin files

• Bin files include jpeg image data pairs, and a label indicating if it's a same person, so there are double images than labels

  # bins | issame_list
  img_1 img_2 | 1
  img_3 img_4 | 1
  img_5 img_6 | 0
  img_7 img_8 | 0

• Image data in bin files like CFP-FP AgeDB-30 is not compatible with tf.image.decode_jpeg, we need to reformat it.
• Python code

  import io
  import pickle
  import tensorflow as tf
  from skimage.io import imread
  
  test_bin_files = ['/datasets/faces_emore/agedb_30.bin', '/datasets/faces_emore/cfp_fp.bin']
  print("test_bin_files =", test_bin_files)
  for test_bin_file in test_bin_files:
      with open(test_bin_file, "rb") as ff:
          bins, issame_list = pickle.load(ff, encoding="bytes")
  
      bb = [bytes(ii) for ii in bins]
      print("Saving to %s" % test_bin_file)
      with open(test_bin_file, "wb") as ff:
          pickle.dump([bb, issame_list], ff)

Conclusion

• After these steps, we made our training dataset folder and evaluating dataset bin files

---

[leondgarse]

Basic Implementation

---

Introduce

• In this basic implementation, we will realize:
  • A basic model
  • Softmax and arcface loss function
  • Training only by training dataset folder

Default import

• These import will be default

  import os
  import sys
  import pandas as pd
  import numpy as np
  import tensorflow as tf
  from tensorflow import keras
  
  gpus = tf.config.experimental.list_physical_devices("GPU")
  for gpu in gpus:
      tf.config.experimental.set_memory_growth(gpu, True)

Loading data by ImageDataGenerator

• ImageDataGenerator is used for reading image data and image augment, but is not very efficiency in speed.
• Python code construct an ImageDataGenerator from data folder. Output is images data and one-hot labels in batch.

  ''' get image paths from data folder '''
  import glob2
  import pickle
  image_names = glob2.glob('/datasets/faces_emore_112x112_folders/*/*.jpg')
  image_names = np.random.permutation(image_names).tolist()
  image_classes = [int(os.path.basename(os.path.dirname(ii))) for ii in image_names]
  
  ''' Backup to pickle '''
  with open('faces_emore_img_class_shuffle.pkl', 'wb') as ff:
      pickle.dump({'image_names': image_names, "image_classes": image_classes}, ff)

  Next time, we can just restore those image paths from pickle.

  ''' Restore from pickle '''
  import pickle
  from keras.preprocessing.image import ImageDataGenerator
  
  AUTOTUNE = tf.data.experimental.AUTOTUNE
  with open('faces_emore_img_class_shuffle.pkl', 'rb') as ff:
      aa = pickle.load(ff)
  image_names, image_classes = aa['image_names'], aa['image_classes']
  print(len(image_names), len(image_classes))
  # 5822653 5822653
  
  ''' Construct a dataframe feed to ImageDataGenerator '''
  data_df = pd.DataFrame({"image_names": image_names, "image_classes": image_classes})
  data_df.image_classes = data_df.image_classes.map(str)
  
  ''' ImageDataGenerator flow_from_dataframe '''
  batch_size = 128
  image_gen = ImageDataGenerator(rescale=1./255, horizontal_flip=True, validation_split=0.05)
  train_data_gen = image_gen.flow_from_dataframe(data_df, directory=None, x_col='image_names', y_col="image_classes", class_mode='categorical', target_size=(112, 112), batch_size=batch_size, subset='training', validate_filenames=False)
  # Found 5240388 non-validated image filenames belonging to 85742 classes.
  val_data_gen = image_gen.flow_from_dataframe(data_df, directory=None, x_col='image_names', y_col="image_classes", class_mode='categorical', target_size=(112, 112), batch_size=batch_size, subset='validation', validate_filenames=False)
  # Found 582265 non-validated image filenames belonging to 85742 classes.
  
  classes = data_df.image_classes.unique().shape[0]
  steps_per_epoch = np.ceil(len(train_data_gen.classes) / batch_size)
  validation_steps = np.ceil(len(val_data_gen.classes) / batch_size)

  The training data is split into two datasets, train_data_gen and val_data_gen, we will just use val_data_gen for evaluation.

Convert ImageDataGenerator to dataset

• tf.data.Dataset can speed up data reading significantly

  ''' Convert to tf.data.Dataset '''
  train_ds = tf.data.Dataset.from_generator(lambda: train_data_gen, output_types=(tf.float32, tf.int32), output_shapes=([None, 112, 112, 3], [None, classes]))
  # train_ds = train_ds.cache()
  train_ds = train_ds.repeat()
  train_ds = train_ds.prefetch(buffer_size=AUTOTUNE)
  
  val_ds = tf.data.Dataset.from_generator(lambda: val_data_gen, output_types=(tf.float32, tf.int32), output_shapes=([None, 112, 112, 3], [None, classes]))
  
  xx, yy = next(iter(train_ds))
  print(xx.shape, yy.shape)
  # (128, 112, 112, 3) (128, 85742)

Basic model

• Our basic model will include:
  • A backbone like ResNet50 or ResNet101 or MobileNet or any others
  • A GlobalAveragePooling2D or Flatten layer
  • An Dense layer outputs embedding values will be used for classification
  • In this basic implementation, we will add different bottleneck for different loss functions
• Python code

  from tensorflow.keras import layers
  
  ''' Basic model '''
  # xx = tf.keras.applications.MobileNetV2(input_shape=(112, 112, 3), include_top=False, weights='imagenet')
  xx = tf.keras.applications.ResNet50V2(input_shape=(112, 112, 3), include_top=False, weights='imagenet')
  xx.trainable = True
  
  inputs = xx.inputs[0]
  nn = xx.outputs[0]
  nn = layers.GlobalAveragePooling2D()(nn)
  nn = layers.Dropout(0.1)(nn)
  embedding = layers.Dense(512, name='embedding')(nn)
  basic_model = keras.models.Model(inputs, embedding)
  basic_model.summary()

Softmax loss train

• This could be the simplest implementation for this face classification mission
  • Softmax train could set a baseline for our model, and the convergence speed is faster than other losses, so we could use this basic model for other loss training
  • Just notice if the loss is decreasing and accuracy is increasing, if this basic training is not working this way, then there must be something wrong
• Python code we only add a softmax layer on top of the basic model

  ''' Model definition '''
  output = layers.Dense(classes, activation='softmax')(basic_model.outputs[0])
  model = keras.models.Model(basic_model.inputs[0], output)
  model.summary()
  
  ''' Compile and fit '''
  model.compile(optimizer='adamax', loss='categorical_crossentropy', metrics=["accuracy"])
  hist = model.fit(train_ds, epochs=20, verbose=1, steps_per_epoch=steps_per_epoch, validation_data=val_ds, validation_steps=validation_steps)

• Training process just run a few step

  Epoch 1/20
  43216/43216 [==============================] - 8650s 200ms/step - loss: 3.8481 - accuracy: 0.4496 - val_loss: 2.6660 - val_accuracy: 0.5180
  Epoch 2/20
  43216/43216 [==============================] - 8792s 203ms/step - loss: 0.9634 - accuracy: 0.8118 - val_loss: 1.2425 - val_accuracy: 0.7599
  Epoch 3/20
  43216/43216 [==============================] - 8720s 202ms/step - loss: 0.6660 - accuracy: 0.8676 - val_loss: 1.3942 - val_accuracy: 0.7380
  Epoch 4/20
  43216/43216 [==============================] - 8713s 202ms/step - loss: 0.5394 - accuracy: 0.8920 - val_loss: 0.6720 - val_accuracy: 0.8733
  Epoch 5/20
  43216/43216 [==============================] - 8873s 205ms/step - loss: 0.4662 - accuracy: 0.9063 - val_loss: 0.7837 - val_accuracy: 0.8540

Arcface loss train

• We have some realization only differ in detail, and there is no big difference in result
• Model definition we add a WeightNormalization layer on top of basic model
  • The first realization is only normalize weight in this layer, and normalize embedding in loss function
  • The second realization is normalize both weight and embedding in this layer, we will use this in our basic implementation
  • In custom keras layer inheriting keras.layers.Layer, other function like get_config should also be implemented for model serialization

  ''' WeightNormalization layer <second realization> '''
  class NormDense(keras.layers.Layer):
      def __init__(self, units=1000, **kwargs):
          super(NormDense, self).__init__(**kwargs)
          self.init = keras.initializers.glorot_normal()
          self.units = units
  
      def build(self, input_shape):
          self.w = self.add_weight(name="norm_dense_w", shape=(input_shape[-1], self.units), initializer=self.init, trainable=True)
          super(NormDense, self).build(input_shape)
  
      def call(self, inputs, **kwargs):
          norm_w = tf.nn.l2_normalize(self.w, axis=0)
          norm_inputs = tf.nn.l2_normalize(inputs, axis=1)
          return tf.matmul(norm_inputs, norm_w)
  
      def compute_output_shape(self, input_shape):
          return (input_shape[0], self.units)
  
      def get_config(self):
          config = super(NormDense, self).get_config()
          config.update({"units": self.units})
          return config
  
      @classmethod
      def from_config(cls, config):
          return cls(**config)
  
  ''' Model definition '''
  inputs = basic_model.inputs[0]
  embedding = basic_model.outputs[0]
  output = NormDense(classes, name='norm_dense')(embedding)
  model = keras.models.Model(inputs, output)
  model.summary()

• Arcface loss function

  def arcface_loss(y_true, y_pred, margin1=1.0, margin2=0.5, margin3=0.0, scale=64.0, from_logits=True, label_smoothing=0):
      threshold = np.cos((np.pi - margin2) / margin1)
      theta_min = (-1 - margin3) * 2
      norm_logits = y_pred
      y_pred_vals = norm_logits[tf.cast(y_true, dtype=tf.bool)]
      # y_pred_vals = tf.clip_by_value(y_pred_vals, clip_value_min=-1.0, clip_value_max=1.0)
      if margin1 == 1.0 and margin3 == 0.0:
          theta = tf.cos(tf.acos(y_pred_vals) + margin2)
      else:
          theta = tf.cos(tf.acos(y_pred_vals) * margin1 + margin2) - margin3
      theta_valid = tf.where(y_pred_vals > threshold, theta, theta_min - theta)
      theta_one_hot = tf.expand_dims(theta_valid - y_pred_vals, 1) * tf.cast(y_true, dtype=tf.float32)
      arcface_logits = (theta_one_hot + norm_logits) * scale
      # tf.assert_equal(tf.math.is_nan(tf.reduce_mean(arcface_logits)), False)
      return tf.keras.losses.categorical_crossentropy(
          y_true, arcface_logits, from_logits=from_logits, label_smoothing=label_smoothing
      )

• Model compile and fit

  model.compile(optimizer='adamax', loss=arcface_loss, metrics=["accuracy"])
  hist = model.fit(train_ds, epochs=20, verbose=1, steps_per_epoch=steps_per_epoch, validation_data=val_ds, validation_steps=validation_steps)

• Training process just run a few step. We can see the loss is decreasing, but the accuracy is increasing too slow

  Epoch 1/20
  43216/43216 [==============================] - 9148s 212ms/step - loss: 34.7171 - accuracy: 0.0088 - val_loss: 14.4753 - val_accuracy: 0.0010
  Epoch 2/20
  43216/43216 [==============================] - 8995s 208ms/step - loss: 14.4870 - accuracy: 0.0022 - val_loss: 14.2573 - val_accuracy: 0.0118
  Epoch 3/20
  43216/43216 [==============================] - 8966s 207ms/step - loss: 14.5741 - accuracy: 0.0146 - val_loss: 14.6334 - val_accuracy: 0.0156
  Epoch 4/20
  43216/43216 [==============================] - 8939s 207ms/step - loss: 14.6519 - accuracy: 0.0175 - val_loss: 14.3232 - val_accuracy: 0.0158
  Epoch 5/20
  43216/43216 [==============================] - 9122s 211ms/step - loss: 14.6973 - accuracy: 0.0198 - val_loss: 15.0081 - val_accuracy: 0.0210

Combined train

• We use the same basic model for softmax and arcface train, so we can use a trained basic model by softmax, and fine-tune by acrface loss

  ''' Save basic model weights in softmax loss train '''
  basic_model.save_weights('basic_model_softmax.h5')
  
  ''' Load basic model weights before arcface train fit '''
  basic_model.load_weights('basic_model_softmax.h5')
  
  ''' Train for a few epoch only to fit the bottleneck layer '''
  basic_model.trainable = False
  model.compile(optimizer='adamax', loss=arcface_loss, metrics=["accuracy"])
  hist = model.fit(train_ds, epochs=2, verbose=1, steps_per_epoch=steps_per_epoch, validation_data=val_ds, validation_steps=validation_steps)
  
  ''' Train the whole model '''
  basic_model.trainable = True
  model.compile(optimizer='adamax', loss=arcface_loss, metrics=["accuracy"])  # MUST run compile after changing trainable value
  hist = model.fit(train_ds, epochs=20, verbose=1, steps_per_epoch=steps_per_epoch, validation_data=val_ds, validation_steps=validation_steps)

---

[leondgarse]

Advanced Implementation

---

Introduce

• In this advanced implementation, we will realize:
  • Loading data by Datasets directly.
  • Evaluating callbacks using bin files.
  • Learning rate and model checkpoint callbacks.

Default import

• These import will be default

  import os
  import sys
  import pandas as pd
  import numpy as np
  import tensorflow as tf
  from tensorflow import keras
  
  gpus = tf.config.experimental.list_physical_devices("GPU")
  for gpu in gpus:
      tf.config.experimental.set_memory_growth(gpu, True)

Loading data by Datasets

• We will use bin files for evaluation, so the whole training data can be used for training

  ''' get image paths from data folder '''
  import glob2
  import pickle
  image_names = glob2.glob('/datasets/faces_emore_112x112_folders/*/*.jpg')
  image_names = np.random.permutation(image_names).tolist()
  image_classes = [int(os.path.basename(os.path.dirname(ii))) for ii in image_names]
  class_dict = {ii: id for id, ii in enumerate(np.unique(image_classes))}
  image_classes_int = [class_dict[ii] for ii in image_classes]
  
  ''' Backup to pickle '''
  with open('faces_emore_img_class_shuffle.pkl', 'wb') as ff:
      pickle.dump({'image_names': image_names, "image_classes": image_classes_int}, ff)

  Next time, we can just restore those image paths from pickle.

  ''' Restore from pickle '''
  import pickle
  AUTOTUNE = tf.data.experimental.AUTOTUNE
  with open('faces_emore_img_class_shuffle.pkl', 'rb') as ff:
      aa = pickle.load(ff)
  image_names, image_classes = aa['image_names'], aa['image_classes']
  classes = np.max(image_classes) + 1
  print(len(image_names), len(image_classes), classes)
  # 5822653 5822653 85742
  
  ''' Construct a dataset by image_names '''
  # list_ds = tf.data.Dataset.list_files('/datasets/faces_emore_112x112_folders/*/*')
  list_ds = tf.data.Dataset.from_tensor_slices((image_names, image_classes))
  
  ''' Construct a dataset generating images and one-hot labels '''
  def process_path(file_path, label, classes, img_shape=(112, 112)):
      label = tf.one_hot(label, depth=classes, dtype=tf.int32)
      img = tf.io.read_file(file_path)
      img = tf.image.decode_jpeg(img, channels=3)
      img = tf.image.convert_image_dtype(img, tf.float32)
      img = tf.image.resize(img, img_shape)
      img = tf.image.random_flip_left_right(img)
      return img, label
  
  def prepare_for_training(ds, classes, cache=False, shuffle_buffer_size=None, batch_size=128):
      if cache:
          ds = ds.cache(cache) if isinstance(cache, str) else ds.cache()
      if shuffle_buffer_size == None:
          shuffle_buffer_size = batch_size * 100
  
      ds = ds.shuffle(buffer_size=shuffle_buffer_size)
      ds = ds.repeat()
      ds = ds.map(lambda xx, yy: process_path(xx, yy, classes), num_parallel_calls=AUTOTUNE)
      ds = ds.batch(batch_size)
      ds = ds.prefetch(buffer_size=AUTOTUNE)
      return ds
  
  # train_ds = prepare_for_training(labeled_ds, cache="/tmp/faces_emore.tfcache")
  batch_size = 128 * len(tf.config.experimental.get_visible_devices('GPU'))
  train_ds = prepare_for_training(list_ds, classes, cache=False, shuffle_buffer_size=len(image_names), batch_size=batch_size)
  steps_per_epoch = np.ceil(len(image_names) / batch_size)
  
  image_batch, label_batch = next(iter(train_ds))
  print(image_batch.shape, label_batch.shape)
  # (128, 112, 112, 3) (128, 85742)

Evaluation callback

• Evaluation is triggered on epoch end, so we can use custom callbacks to realize this
• Bin files includes image data bins and is same labels. We can measure if two images is same person or not by their distance
  • Distance can be measured by cosine or pairwise or others. Here we will use cosine distance.
  • threshold is a number, that any distance greater (cosine distance) / less (pairwise distance) than this value is judged as a same person
• Python code implement a sub class of tf.keras.callbacks.Callback, re-write on_epoch_end function to do evaluation, and can also save model

  import pickle
  import io
  from tqdm import tqdm
  from skimage.io import imread
  from sklearn.preprocessing import normalize
  
  class epoch_eval_callback(tf.keras.callbacks.Callback):
      def __init__(self, basic_model, test_bin_file, batch_size=128, rescale=1./255, save_model=None):
          super(epoch_eval_callback, self).__init__()
          self.basic_model = basic_model
          bins, issame_list = np.load(test_bin_file, encoding='bytes', allow_pickle=True)
          ds = tf.data.Dataset.from_tensor_slices(bins)
          _imread = lambda xx: tf.image.convert_image_dtype(tf.image.decode_jpeg(xx), dtype=tf.float32)
          ds = ds.map(_imread)
          self.ds = ds.batch(128)
          self.test_issame = np.array(issame_list)
          self.test_names = os.path.splitext(os.path.basename(test_bin_file))[0]
          self.max_accuracy = 0
          self.batch_size = batch_size
          self.steps = int(np.ceil(len(bins) / batch_size))
          self.save_model = save_model
  
      # def on_batch_end(self, batch=0, logs=None):
      def on_epoch_end(self, epoch=0, logs=None):
          dists = []
          embs = []
          tf.print("\n")
          for img_batch in tqdm(self.ds, 'Evaluating ' + self.test_names, total=self.steps):
              emb = self.basic_model.predict(img_batch)
              embs.extend(emb)
          embs = np.array(embs)
          if np.isnan(embs).sum() != 0:
              tf.print("NAN in embs, not a good one")
              return
          embs = normalize(embs)
          embs_a = embs[::2]
          embs_b = embs[1::2]
          dists = (embs_a * embs_b).sum(1)
  
          self.tt = np.sort(dists[self.test_issame[:dists.shape[0]]])
          self.ff = np.sort(dists[np.logical_not(self.test_issame[:dists.shape[0]])])
  
          max_accuracy = 0
          thresh = 0
          for vv in reversed(self.ff[-300:]):
              acc_count = (self.tt > vv).sum() + (self.ff <= vv).sum()
              acc = acc_count / dists.shape[0]
              if acc > max_accuracy:
                  max_accuracy = acc
                  thresh = vv
          tf.print("\n")
          if max_accuracy > self.max_accuracy:
              is_improved = True
              self.max_accuracy = max_accuracy
              if self.save_model:
                  save_path = '%s_%d' % (self.save_model, epoch)
                  tf.print("Saving model to: %s" % (save_path))
                  model.save(save_path)
          else:
              is_improved = False
          tf.print(">>>> %s evaluation max accuracy: %f, thresh: %f, overall max accuracy: %f, improved = %s" % (self.test_names, max_accuracy, thresh, self.max_accuracy, is_improved))

• Test We can use the basic model or any other model outputs embeddings to test this function

  basic_model = ...
  basic_model.load_weights('basic_model_softmax.h5')
  
  aa = epoch_eval_callback(basic_model, '/datasets/faces_emore/lfw.bin', save_model='./test')
  aa.on_epoch_end()

Callbacks

• Callbacks includes
  • ModelCheckpoint saves model weights on every epoch end
  • ReduceLROnPlateau adjusts learning rate on every epoch end
  • Evaluation functions evaluate model accuracy on every epoch end
  • Evaluation functions can also be set to saving model weights on epoch end according to the evaluating accuracy
• Python code

  ''' Callbacks '''
  from tensorflow.keras.callbacks import ModelCheckpoint, ReduceLROnPlateau
  
  # lfw_eval = epoch_eval_callback('/datasets/faces_emore/lfw.bin', save_model="keras_model_lfw_acc.h5")
  lfw_eval = epoch_eval_callback('/datasets/faces_emore/lfw.bin', save_model=None)
  agedb_30_eval = epoch_eval_callback('/datasets/faces_emore/agedb_30.bin', save_model=None)
  cfp_fp_eval = epoch_eval_callback('/datasets/faces_emore/cfp_fp.bin', save_model=None)
  
  # reduce_lr = ReduceLROnPlateau('val_loss', factor=0.1, patience=5, verbose=1)
  def scheduler(epoch):
      lr = 0.001 if epoch < 10 else 0.001 * np.exp(0.2 * (10 - epoch))
      print('\nLearning rate for epoch {} is {}'.format(epoch + 1, lr))
      return lr
  lr_scheduler = tf.keras.callbacks.LearningRateScheduler(scheduler)
  
  model_checkpoint = ModelCheckpoint("./keras_checkpoints.h5", verbose=1)
  # model_checkpoint = ModelCheckpoint("./keras_checkpoints_res_arcface", 'val_loss', verbose=1, save_best_only=True)
  callbacks = [lr_scheduler, lfw_eval, agedb_30_eval, cfp_fp_eval, model_checkpoint]

Softmax loss train

from tensorflow import keras

''' Build model '''
xx = tf.keras.applications.MobileNetV2(input_shape=(112, 112, 3), include_top=False, weights='imagenet')
xx.trainable = True

inputs = xx.inputs[0]
nn = xx.outputs[0]
nn = keras.layers.GlobalAveragePooling2D()(nn)
# nn = keras.layers.Dropout(0.1)(nn)
embedding = keras.layers.Dense(512, name='embedding')(nn)
basic_model = keras.models.Model(inputs, embedding)

output = keras.layers.Dense(classes, activation='softmax')(basic_model.outputs[0])
model = keras.models.Model(basic_model.inputs[0], output)
# model.summary()

model.compile(optimizer='adamax', loss=keras.losses.CategoricalCrossentropy(label_smoothing=0.1), metrics="acc")
hist = model.fit(train_ds, epochs=20, verbose=1, callbacks=callbacks, steps_per_epoch=steps_per_epoch, initial_epoch=0)