{"cells":[{"metadata":{},"cell_type":"markdown","source":"### Work in progress/incomplete\n\n#### Feel free to give feedback!","execution_count":null},{"metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","trusted":true,"_kg_hide-input":true},"cell_type":"code","source":"import numpy as np\nimport tensorflow as tf\nimport math\nimport pandas as pd\nfrom sklearn import model_selection\nimport glob\nimport os\nfrom zipfile import ZipFile\nimport shutil\nimport tqdm.notebook as tqdm\n\nimport logging\ntf.get_logger().setLevel(logging.ERROR)\nimport warnings\nwarnings.filterwarnings(\"ignore\")\n\ngpus = tf.config.experimental.list_physical_devices('GPU')\nnum_gpus = len(gpus)\nmixed_precision = False\nif gpus:\n    try:\n        for gpu in gpus:\n            tf.config.experimental.set_memory_growth(gpu, True)\n        logical_gpus = tf.config.experimental.list_logical_devices('GPU')\n        print(num_gpus, \"Physical GPUs,\", len(logical_gpus), \"Logical GPUs\")\n    except RuntimeError as e:\n        print(e)\n\n    # policy = tf.keras.mixed_precision.experimental.Policy('mixed_float16')\n    # tf.keras.mixed_precision.experimental.set_policy(policy)\n    # print('Compute dtype: %s' % policy.compute_dtype)\n    # print('Variable dtype: %s' % policy.variable_dtype)\n    # mixed_precision = True\n\n    \nif num_gpus == 0:\n    strategy = tf.distribute.OneDeviceStrategy(device='CPU')\n    print(\"Setting strategy to OneDeviceStrategy(device='CPU')\")\nelif num_gpus == 1:\n    strategy = tf.distribute.OneDeviceStrategy(device='GPU')\n    print(\"Setting strategy to OneDeviceStrategy(device='GPU')\")\nelse:\n    strategy = tf.distribute.MirroredStrategy()\n    print(\"Setting strategy to MirroredStrategy()\")","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"config = {\n    'learning_rate': 5e-3,\n    'momentum': 0.9,\n    'scale': 30,\n    'margin': 0.3,\n    'clip_grad': 10.0,\n    'n_epochs': 4,\n    'batch_size': 32,\n    'input_size': (384, 384, 3),\n    'n_classes': 81313,\n    'dense_units': 512,\n    'dropout_rate': 0.0,\n}","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"def read_df(input_path, alpha=0.5):\n    files_paths = glob.glob(input_path + 'train/*/*/*/*')\n    mapping = {}\n    for path in files_paths:\n        mapping[path.split('/')[-1].split('.')[0]] = path\n    df = pd.read_csv(input_path + 'train.csv')\n    df['path'] = df['id'].map(mapping)\n    \n    counts_map = dict(\n        df.groupby('landmark_id')['path'].agg(lambda x: len(x)))\n    df['counts'] = df['landmark_id'].map(counts_map)\n    df['probs'] = (\n        (1/df.counts**alpha) / (1/df.counts**alpha).max()).astype(np.float32)\n    uniques = df['landmark_id'].unique()\n    uniques_map = dict(zip(uniques, range(len(uniques))))\n    df['labels'] = df['landmark_id'].map(uniques_map)\n    return df\n\ndf = read_df('../input/landmark-retrieval-2020/')\ndf.head(10)","execution_count":null,"outputs":[]},{"metadata":{"_uuid":"d629ff2d2480ee46fbb7e2d37f6b5fab8052498a","_cell_guid":"79c7e3d0-c299-4dcb-8224-4455121ee9b0","trusted":true},"cell_type":"code","source":"def _get_transform_matrix(rotation, shear, hzoom, wzoom, hshift, wshift):\n\n    def get_3x3_mat(lst):\n        return tf.reshape(tf.concat([lst],axis=0), [3,3])\n\n    # convert degrees to radians\n    rotation = math.pi * rotation / 360.\n    shear    = math.pi * shear    / 360.\n\n    one  = tf.constant([1],dtype='float32')\n    zero = tf.constant([0],dtype='float32')\n\n    c1   = tf.math.cos(rotation)\n    s1   = tf.math.sin(rotation)\n    rot_mat = get_3x3_mat([c1,    s1,   zero ,\n                           -s1,   c1,   zero ,\n                           zero,  zero, one ])\n\n    c2 = tf.math.cos(shear)\n    s2 = tf.math.sin(shear)\n    shear_mat = get_3x3_mat([one,  s2,   zero ,\n                             zero, c2,   zero ,\n                             zero, zero, one ])\n\n    zoom_mat = get_3x3_mat([one/hzoom, zero,      zero,\n                            zero,      one/wzoom, zero,\n                            zero,      zero,      one])\n\n    shift_mat = get_3x3_mat([one,  zero, hshift,\n                             zero, one,  wshift,\n                             zero, zero, one   ])\n\n    return tf.matmul(\n        tf.matmul(rot_mat, shear_mat),\n        tf.matmul(zoom_mat, shift_mat)\n    )\n\ndef _spatial_transform(image,\n                       rotation=3.0,\n                       shear=2.0,\n                       hzoom=8.0,\n                       wzoom=8.0,\n                       hshift=8.0,\n                       wshift=8.0):\n\n    ydim = tf.gather(tf.shape(image), 0)\n    xdim = tf.gather(tf.shape(image), 1)\n    xxdim = xdim % 2\n    yxdim = ydim % 2\n\n    # random rotation, shear, zoom and shift\n    rotation = rotation * tf.random.normal([1], dtype='float32')\n    shear = shear * tf.random.normal([1], dtype='float32')\n    hzoom = 1.0 + tf.random.normal([1], dtype='float32') / hzoom\n    wzoom = 1.0 + tf.random.normal([1], dtype='float32') / wzoom\n    hshift = hshift * tf.random.normal([1], dtype='float32')\n    wshift = wshift * tf.random.normal([1], dtype='float32')\n\n    m = _get_transform_matrix(\n        rotation, shear, hzoom, wzoom, hshift, wshift)\n\n    # origin pixels\n    y = tf.repeat(tf.range(ydim//2, -ydim//2,-1), xdim)\n    x = tf.tile(tf.range(-xdim//2, xdim//2), [ydim])\n    z = tf.ones([ydim*xdim], dtype='int32')\n    idx = tf.stack([y, x, z])\n\n    # destination pixels\n    idx2 = tf.matmul(m, tf.cast(idx, dtype='float32'))\n    idx2 = tf.cast(idx2, dtype='int32')\n    # clip to origin pixels range\n    idx2y = tf.clip_by_value(idx2[0,], -ydim//2+yxdim+1, ydim//2)\n    idx2x = tf.clip_by_value(idx2[1,], -xdim//2+xxdim+1, xdim//2)\n    idx2 = tf.stack([idx2y, idx2x, idx2[2,]])\n\n    # apply destinations pixels to image\n    idx3 = tf.stack([ydim//2-idx2[0,], xdim//2-1+idx2[1,]])\n    d = tf.gather_nd(image, tf.transpose(idx3))\n    image = tf.reshape(d, [ydim, xdim, 3])\n    return image\n\ndef _pixel_transform(image,\n                     saturation_delta=0.3,\n                     contrast_delta=0.1,\n                     brightness_delta=0.2):\n    image = tf.image.random_saturation(\n        image, 1-saturation_delta, 1+saturation_delta)\n    image = tf.image.random_contrast(\n        image, 1-contrast_delta, 1+contrast_delta)\n    image = tf.image.random_brightness(\n        image, brightness_delta)\n    return image\n\ndef _random_fliplr(image, p=0.25):\n    r = tf.random.uniform(())\n    mirror_cond = tf.math.less(r, p)\n    image = tf.cond(\n        mirror_cond,\n        lambda: tf.reverse(image, [1]),\n        lambda: image\n    )\n    return image\n\ndef preprocess_input(image, target_size, augment=False):\n    \n    image = tf.image.resize(\n        image, target_size, method='bilinear')\n\n    image = tf.cast(image, tf.uint8)\n    if augment:\n        image = _spatial_transform(image)\n        image = _random_fliplr(image)\n        image = _pixel_transform(image)\n    image = tf.cast(image, tf.float32)\n    image /= 255.\n    return image\n\ndef create_dataset(df, training, batch_size, input_size):\n\n    def read_image(image_path):\n        image = tf.io.read_file(image_path)\n        return tf.image.decode_jpeg(image, channels=3)\n    \n    def filter_by_probs(x, y, p):\n        if p > np.random.uniform(0, 1):\n            return True\n        return False\n\n    image_paths, labels, probs = df.path, df.labels, df.probs\n\n    dataset = tf.data.Dataset.from_tensor_slices((image_paths, labels, probs))\n    if training:\n        dataset = dataset.shuffle(100_000)\n    dataset = dataset.map(\n        lambda x, y, p: (read_image(x), y, p),\n        tf.data.experimental.AUTOTUNE)\n    if training:\n        dataset = dataset.filter(filter_by_probs)\n    dataset = dataset.map(\n        lambda x, y, p: (preprocess_input(x, input_size[:2], training), y),\n        tf.data.experimental.AUTOTUNE)\n    dataset = dataset.batch(batch_size)\n    return dataset","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"class AddMarginProduct(tf.keras.layers.Layer):\n    '''\n    Implements large margin cosine distance.\n\n    References:\n        https://arxiv.org/pdf/1801.07698.pdf\n        https://github.com/lyakaap/Landmark2019-1st-and-3rd-Place-Solution/\n            blob/master/src/modeling/metric_learning.py\n    '''\n    def __init__(self, n_classes, s=30, m=0.30, **kwargs):\n        super(AddMarginProduct, self).__init__(**kwargs)\n        self.n_classes = n_classes\n        self.s = s\n        self.m = m\n\n    def build(self, input_shape):\n        super(AddMarginProduct, self).build(input_shape[0])\n\n        self.W = self.add_weight(\n            name='W',\n            shape=(int(input_shape[0][-1]), self.n_classes),\n            initializer='glorot_uniform',\n            dtype='float32',\n            trainable=True,\n            regularizer=None)\n\n    def call(self, inputs):\n        X, y = inputs\n        y = tf.cast(y, dtype=tf.int32)\n        cosine = tf.matmul(\n            tf.math.l2_normalize(X, axis=1),\n            tf.math.l2_normalize(self.W, axis=0)\n        )\n        phi = cosine - self.m\n        one_hot = tf.cast(\n            tf.one_hot(y, depth=self.n_classes),\n            dtype=cosine.dtype\n        )\n        output = (one_hot * phi) + ((1.0 - one_hot) * cosine)\n        output *= self.s\n        return output\n\n\ndef create_model(input_shape,\n                 n_classes,\n                 dense_units=512,\n                 dropout_rate=0.0,\n                 scale=30,\n                 margin=0.3):\n\n    backbone = tf.keras.applications.ResNet50(\n        include_top=False,\n        input_shape=input_shape,\n        weights=('../input/imagenet-weights/' +\n                 'resnet50_weights_tf_dim_ordering_tf_kernels_notop.h5')\n    )\n\n    pooling = tf.keras.layers.GlobalAveragePooling2D(name='head/pooling')\n    dropout = tf.keras.layers.Dropout(dropout_rate, name='head/dropout')\n    dense = tf.keras.layers.Dense(dense_units, name='head/dense')\n\n    margin = AddMarginProduct(\n        n_classes=n_classes,\n        s=scale,\n        m=margin,\n        name='head/cos_margin',\n        dtype='float32')\n\n    softmax = tf.keras.layers.Softmax(dtype='float32')\n\n    image = tf.keras.layers.Input(input_shape, name='input/image')\n    label = tf.keras.layers.Input((), name='input/label')\n\n    x = backbone(image)\n    x = pooling(x)\n    # x = dropout(x)\n    # x = dense(x)\n    x = margin([x, label])\n    x = softmax(x)\n    return tf.keras.Model(\n        inputs=[image, label], outputs=x)\n\n\nclass DistributedModel:\n\n    def __init__(self,\n                 input_size,\n                 n_classes,\n                 batch_size,\n                 finetuned_weights,\n                 dense_units,\n                 dropout_rate,\n                 scale,\n                 margin,\n                 optimizer,\n                 strategy,\n                 mixed_precision,\n                 clip_grad):\n\n        self.model = create_model(\n            input_shape=input_size,\n            n_classes=n_classes,\n            dense_units=dense_units,\n            dropout_rate=dropout_rate,\n            scale=scale,\n            margin=margin,)\n\n        self.input_size = input_size\n        self.global_batch_size = batch_size * strategy.num_replicas_in_sync\n\n        if finetuned_weights:\n            self.model.load_weights(finetuned_weights)\n\n        self.mixed_precision = mixed_precision\n        self.optimizer = optimizer\n        self.strategy = strategy\n        self.clip_grad = clip_grad\n\n        # loss function\n        self.loss_object = tf.keras.losses.SparseCategoricalCrossentropy(\n            from_logits=False, reduction=tf.keras.losses.Reduction.NONE)\n\n        # metrics\n        self.mean_loss_train = tf.keras.metrics.SparseCategoricalCrossentropy(\n            from_logits=False)\n        self.mean_accuracy_train = tf.keras.metrics.SparseTopKCategoricalAccuracy(\n            k=5)\n\n        if self.optimizer and self.mixed_precision:\n            self.optimizer = \\\n                tf.keras.mixed_precision.experimental.LossScaleOptimizer(\n                    optimizer, loss_scale='dynamic')\n\n    def _compute_loss(self, labels, probs):\n        per_example_loss = self.loss_object(labels, probs)\n        return tf.nn.compute_average_loss(\n            per_example_loss, global_batch_size=self.global_batch_size)\n\n    def _backprop_loss(self, tape, loss, weights):\n        gradients = tape.gradient(loss, weights)\n        if self.mixed_precision:\n            gradients = self.optimizer.get_unscaled_gradients(gradients)\n        clipped, _ = tf.clip_by_global_norm(gradients, clip_norm=self.clip_grad)\n        self.optimizer.apply_gradients(zip(clipped, weights))\n\n    def _train_step(self, inputs):\n        with tf.GradientTape() as tape:\n            probs = self.model(inputs, training=True)\n            loss = self._compute_loss(inputs[1], probs)\n            if self.mixed_precision:\n                loss = self.optimizer.get_scaled_loss(loss)\n        self._backprop_loss(tape, loss, self.model.trainable_weights)\n        self.mean_loss_train.update_state(inputs[1], probs)\n        self.mean_accuracy_train.update_state(inputs[1], probs)\n        return loss\n\n    @tf.function\n    def _distributed_train_step(self, dist_inputs):\n        per_replica_loss = self.strategy.run(self._train_step, args=(dist_inputs,))\n        return self.strategy.reduce(\n            tf.distribute.ReduceOp.SUM, per_replica_loss, axis=None)\n\n    def train(self, train_ds, epochs, save_path):\n        for epoch in range(epochs):\n            dist_train_ds = self.strategy.experimental_distribute_dataset(train_ds)\n            dist_train_ds = tqdm.tqdm(dist_train_ds)\n            for i, inputs in enumerate(dist_train_ds):\n                loss = self._distributed_train_step(inputs)\n                dist_train_ds.set_description(\n                    \"TRAIN: Loss {:.3f}, Accuracy {:.3f}\".format(\n                        self.mean_loss_train.result().numpy(),\n                        self.mean_accuracy_train.result().numpy()\n                    )\n                )\n            if save_path:\n                self.model.save_weights(save_path)\n\n            self.mean_loss_train.reset_states()\n            self.mean_accuracy_train.reset_states()\n","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"train_ds = create_dataset(\n    df=df,\n    training=True,\n    batch_size=config['batch_size'],\n    input_size=config['input_size'],\n)\n\nwith strategy.scope():\n\n    optimizer = tf.keras.optimizers.SGD(\n        config['learning_rate'], momentum=config['momentum'])\n\n    dist_model = DistributedModel(\n        input_size=config['input_size'],\n        n_classes=config['n_classes'],\n        batch_size=config['batch_size'],\n        finetuned_weights=None,\n        dense_units=config['dense_units'],\n        dropout_rate=config['dropout_rate'],\n        scale=config['scale'],\n        margin=config['margin'],\n        optimizer=optimizer,\n        strategy=strategy,\n        mixed_precision=mixed_precision,\n        clip_grad=config['clip_grad'])\n\n    dist_model.train(\n        train_ds=train_ds, \n        epochs=config['n_epochs'], \n        save_path='model.h5')","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"newmodel = tf.keras.Model(\n    inputs=dist_model.model.get_layer('input/image').input,\n    outputs=dist_model.model.get_layer('head/dense').output)\n\n\n@tf.function(input_signature=[\n        tf.TensorSpec(\n            shape=[None, None, 3],\n            dtype=tf.uint8,\n            name='input_image')\n    ])\ndef serving(input_image):\n    input_image = preprocess_input(\n        input_image, target_size=config['input_size'][:2])\n    outputs = newmodel(input_image[tf.newaxis])\n    features = tf.math.l2_normalize(outputs[0])\n    return {\n        'global_descriptor': tf.identity(features, name='global_descriptor')\n    }\n\n\ntf.saved_model.save(\n    obj=newmodel,\n    export_dir='model',\n    signatures={'serving_default': serving})\n\n\nfilepaths = []\nfor dirpath, _, filepath in os.walk('model'):\n    for fp in filepath:\n        filepaths.append(os.path.join(dirpath, fp))\n\nwith ZipFile('submission.zip', 'w') as zip:\n    for fp in filepaths:\n        print(fp, '/'.join(fp.split('/')[1:]))\n        zip.write(fp, arcname='/'.join(fp.split('/')[1:]))","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}