{"cells":[{"metadata":{"trusted":true},"cell_type":"code","source":"!pip install ../input/kaggle-efficientnet-repo/efficientnet-1.0.0-py3-none-any.whl","execution_count":null,"outputs":[]},{"metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","trusted":true},"cell_type":"code","source":"import os\nimport numpy as np\nimport pandas as pd\nimport argparse\nimport tensorflow as tf\nfrom tensorflow.keras.optimizers import Adam\nfrom tensorflow.keras.losses import categorical_crossentropy\nfrom tensorflow.keras.metrics import categorical_accuracy, top_k_categorical_accuracy\nfrom kaggle_datasets import KaggleDatasets\n\nfrom tensorflow.keras import layers as L\nimport efficientnet.tfkeras as efn\n\n\ndef normalize(image):\n  # https://github.com/tensorflow/tpu/blob/master/models/official/efficientnet/main.py#L325-L326\n  # https://github.com/tensorflow/tpu/blob/master/models/official/efficientnet/efficientnet_builder.py#L31-L32\n  image -= tf.constant([0.485 * 255, 0.456 * 255, 0.406 * 255])  # RGB\n  image /= tf.constant([0.229 * 255, 0.224 * 255, 0.225 * 255])  # RGB\n  return image\n\n\ndef get_model(input_size, backbone='efficientnet-b0', weights='imagenet', tta=False):\n  print(f'Using backbone {backbone} and weights {weights}')\n  x = L.Input(shape=input_size, name='imgs', dtype='float32')\n  y = normalize(x)\n  if backbone.startswith('efficientnet'):\n    model_fn = getattr(efn, f'EfficientNetB{backbone[-1]}')\n\n  y = model_fn(input_shape=input_size, weights=weights, include_top=False)(y)\n  y = L.GlobalAveragePooling2D()(y)\n  y = L.Dropout(0.2)(y)\n  # 1292 of 1295 are present\n  y = L.Dense(1292, activation='softmax')(y)\n  model = tf.keras.Model(x, y)\n\n  if tta:\n    assert False, 'This does not make sense yet'\n    x_flip = tf.reverse(x, [2])  # 'NHWC'\n    y_tta = tf.add(model(x), model(x_flip)) / 2.0\n    tta_model = tf.keras.Model(x, y_tta)\n    return model, tta_model\n\n  return model\n\n\ndef mixup(img_batch, label_batch, batch_size):\n  # https://github.com/tensorpack/tensorpack/blob/master/examples/ResNet/cifar10-preact18-mixup.py\n  weight = tf.random.uniform([batch_size])\n  x_weight = tf.reshape(weight, [batch_size, 1, 1, 1])\n  y_weight = tf.reshape(weight, [batch_size, 1])\n  index = tf.random.shuffle(tf.range(batch_size, dtype=tf.int32))\n  x1, x2 = img_batch, tf.gather(img_batch, index)\n  img_batch = x1 * x_weight + x2 * (1. - x_weight)\n  y1, y2 = label_batch, tf.gather(label_batch, index)\n  label_batch = y1 * y_weight + y2 * (1. - y_weight)\n  return img_batch, label_batch\n\n\ndef get_strategy():\n  # Detect hardware, return appropriate distribution strategy\n  try:\n    tpu = tf.distribute.cluster_resolver.TPUClusterResolver()  # TPU detection\n    print('Running on TPU ', tpu.cluster_spec().as_dict()['worker'])\n  except ValueError:\n    tpu = None\n\n  if tpu:\n    tf.config.experimental_connect_to_cluster(tpu)\n    tf.tpu.experimental.initialize_tpu_system(tpu)\n    strategy = tf.distribute.experimental.TPUStrategy(tpu)\n  else:\n    strategy = tf.distribute.get_strategy()\n\n  print('REPLICAS: ', strategy.num_replicas_in_sync)\n  return strategy\n\n\ndef one_hot(image, label):\n  label = tf.one_hot(label, 1292)\n  return image, label\n\n\ndef read_tfrecords(example, input_size):\n  features = {\n      'img': tf.io.FixedLenFeature([], tf.string),\n      'image_id': tf.io.FixedLenFeature([], tf.int64),\n      'grapheme_root': tf.io.FixedLenFeature([], tf.int64),\n      'vowel_diacritic': tf.io.FixedLenFeature([], tf.int64),\n      'consonant_diacritic': tf.io.FixedLenFeature([], tf.int64),\n      'unique_tuple': tf.io.FixedLenFeature([], tf.int64),\n  }\n  example = tf.io.parse_single_example(example, features)\n  img = tf.image.decode_image(example['img'])\n  img = tf.reshape(img, input_size + (1, ))\n  img = tf.cast(img, tf.float32)\n  # grayscale -> RGB\n  img = tf.repeat(img, 3, -1)\n\n  # image_id = tf.cast(example['image_id'], tf.int32)\n  # grapheme_root = tf.cast(example['grapheme_root'], tf.int32)\n  # vowel_diacritic = tf.cast(example['vowel_diacritic'], tf.int32)\n  # consonant_diacritic = tf.cast(example['consonant_diacritic'], tf.int32)\n  unique_tuple = tf.cast(example['unique_tuple'], tf.int32)\n  return img, unique_tuple\n\n\ndef main():\n  parser = argparse.ArgumentParser()\n  parser.add_argument('--model_id', type=int, default=0)\n  parser.add_argument('--seed', type=int, default=123)\n  parser.add_argument('--lr', type=float, default=2e-4)\n  parser.add_argument('--input_size', type=str, default='160,256')\n  parser.add_argument('--batch_size', type=int, default=128)\n  parser.add_argument('--epochs', type=int, default=25)\n  parser.add_argument('--backbone', type=str, default='efficientnet-b3')\n  parser.add_argument('--weights', type=str, default='imagenet')\n  args, _ = parser.parse_known_args()\n\n  args.input_size = tuple(int(x) for x in args.input_size.split(','))\n  np.random.seed(args.seed)\n  tf.random.set_seed(args.seed)\n\n  # build the model\n  strategy = get_strategy()\n  with strategy.scope():\n    model = get_model(input_size=args.input_size + (3, ), backbone=args.backbone,\n        weights=args.weights)\n\n  model.compile(optimizer=Adam(lr=args.lr),\n                loss=categorical_crossentropy,\n                metrics=[categorical_accuracy, top_k_categorical_accuracy])\n  # print(model.summary())\n  AUTO = tf.data.experimental.AUTOTUNE\n  # create the training and validation datasets\n  ds_path = KaggleDatasets().get_gcs_path('bengali-tfrecords-v010')\n  train_fns = tf.io.gfile.glob(os.path.join(ds_path, 'records/train*.tfrec'))\n  train_ds = tf.data.TFRecordDataset(train_fns, num_parallel_reads=AUTO)\n  train_ds = train_ds.map(lambda e: read_tfrecords(e, args.input_size), num_parallel_calls=AUTO)\n  train_ds = train_ds.repeat().batch(args.batch_size)\n  train_ds = train_ds.map(one_hot, num_parallel_calls=AUTO)\n  train_ds = train_ds.map(lambda a, b: mixup(a, b, args.batch_size), num_parallel_calls=AUTO)\n\n  val_fns = tf.io.gfile.glob(os.path.join(ds_path, 'records/val*.tfrec'))\n  val_ds = tf.data.TFRecordDataset(val_fns, num_parallel_reads=AUTO)\n  val_ds = val_ds.map(lambda e: read_tfrecords(e, args.input_size), num_parallel_calls=AUTO)\n  val_ds = val_ds.batch(args.batch_size)\n  val_ds = val_ds.map(one_hot, num_parallel_calls=AUTO)\n\n  # train\n  num_train_samples = sum(int(fn.split('_')[2]) for fn in train_fns)\n  # num_val_samples = sum(int(fn.split('_')[2]) for fn in val_fns)\n  steps_per_epoch = num_train_samples // args.batch_size\n  print(f'Training on {num_train_samples} samples. Each epochs requires {steps_per_epoch} steps')\n  h = model.fit(train_ds, steps_per_epoch=steps_per_epoch, epochs=args.epochs, verbose=1,\n      validation_data=val_ds)\n  print(h)\n  weight_fn = 'model-%04d.h5' % args.model_id\n  model.save_weights(weight_fn)\n  print(f'Saved weights to: {weight_fn}')\n\n\nmain()\n","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"","execution_count":null,"outputs":[]}],"metadata":{"kernelspec":{"language":"python","display_name":"Python 3","name":"python3"},"language_info":{"pygments_lexer":"ipython3","nbconvert_exporter":"python","version":"3.6.4","file_extension":".py","codemirror_mode":{"name":"ipython","version":3},"name":"python","mimetype":"text/x-python"}},"nbformat":4,"nbformat_minor":4}