{"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_minor":4,"nbformat":4,"cells":[{"cell_type":"code","source":"%%bash\npip install --quiet --upgrade tensorflow-addons comet-ml","metadata":{"execution":{"iopub.status.busy":"2022-02-22T21:19:26.533631Z","iopub.execute_input":"2022-02-22T21:19:26.534137Z","iopub.status.idle":"2022-02-22T21:19:34.872777Z","shell.execute_reply.started":"2022-02-22T21:19:26.534046Z","shell.execute_reply":"2022-02-22T21:19:34.871726Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"from pathlib import Path\n\ndata_root = \"/kaggle/input/h-and-m-personalized-fashion-recommendations\"\ndata_root = Path(data_root)\nlist(map(str, data_root.glob(\"*\")))","metadata":{"execution":{"iopub.status.busy":"2022-02-22T21:19:34.878599Z","iopub.execute_input":"2022-02-22T21:19:34.879004Z","iopub.status.idle":"2022-02-22T21:19:34.892708Z","shell.execute_reply.started":"2022-02-22T21:19:34.878955Z","shell.execute_reply":"2022-02-22T21:19:34.891774Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"import pandas as pd\n\narticles = pd.read_csv(data_root / \"articles.csv\")\nlabel_names = articles.set_index(\"product_type_name\")\nlabel_names = label_names[\"product_type_no\"].to_dict()\nlabel_names = sorted(label_names, key=label_names.get)\narticles = {row[\"article_id\"]: row for row in articles.to_dict(orient=\"records\")}","metadata":{"execution":{"iopub.status.busy":"2022-02-22T21:19:34.895464Z","iopub.execute_input":"2022-02-22T21:19:34.895767Z","iopub.status.idle":"2022-02-22T21:19:39.419091Z","shell.execute_reply.started":"2022-02-22T21:19:34.895740Z","shell.execute_reply":"2022-02-22T21:19:39.418066Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"import tensorflow as tf\nimport tensorflow_addons as tfa\n\ntf.random.set_seed(42)\n\ndef build_decoder(latent_dim):\n    return tf.keras.Sequential([\n        tf.keras.layers.Input([latent_dim]),\n        tf.keras.layers.Dense(len(label_names))\n    ])\n\ndef build_encoder(latent_dim):\n    inputs = tf.keras.layers.Input([None, None, 3])\n    spine = tf.keras.applications.MobileNetV3Small(include_top=False)\n    spine.trainable = False\n    \n    batch = inputs\n    for augment in [\n        tf.keras.layers.Resizing(64, 64),\n        tf.keras.layers.RandomFlip(\"horizontal\"),\n        tf.keras.layers.RandomContrast(0.25),\n        tf.keras.layers.RandomRotation(0.0625),\n        tf.keras.layers.RandomTranslation(0.125, 0.125),\n        tf.keras.layers.RandomZoom(-0.125),\n    ]:\n        batch = augment(batch)\n    vectors = spine(batch, training=False)\n    vectors = tf.keras.layers.GlobalAvgPool2D()(vectors)\n    \n    dims = [512, 256, 128]\n    for i, d in enumerate(dims):\n        name = None if i != range(len(dims))[-1] else \"projection_head\"\n        vectors = tf.keras.layers.Activation(tf.nn.silu)(vectors)\n        vectors = tf.keras.layers.Dropout(0.1)(vectors)\n        vectors = tf.keras.layers.Dense(d, name=name)(vectors)\n        \n    return tf.keras.Model(inputs, vectors)\n\ndef build_optimizer(steps, min_lr, max_lr):\n    sched = tfa.optimizers.TriangularCyclicalLearningRate(\n        min_lr, max_lr, steps / 2\n    )\n    optim = tfa.optimizers.AdaBelief(sched, rectify=True)\n    optim = tfa.optimizers.Lookahead(optim)\n    return optim","metadata":{"execution":{"iopub.status.busy":"2022-02-22T21:19:39.421994Z","iopub.execute_input":"2022-02-22T21:19:39.422284Z","iopub.status.idle":"2022-02-22T21:19:41.733365Z","shell.execute_reply.started":"2022-02-22T21:19:39.422254Z","shell.execute_reply":"2022-02-22T21:19:41.732460Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def unitwise_norm(x):\n    if len(x.get_shape()) <= 1:  # Scalars and vectors\n        axis = None\n        keepdims = False\n    elif len(x.get_shape()) in [2, 3]:  # Linear layers of shape IO or multihead linear\n        axis = 0\n        keepdims = True\n    elif len(x.get_shape()) == 4:  # Conv kernels of shape HWIO\n        axis = [0, 1, 2,]\n        keepdims = True\n    else:\n        raise ValueError(f\"Got a parameter with shape not in [1, 2, 4]! {x}\")\n    return tf.linalg.norm(x, axis=axis, keepdims=keepdims)\n\n\ndef adaptive_clip_grad(parameters, gradients, clip_factor=0.01,\n                       eps=1e-3):\n    new_grads = []\n    for (params, grads) in zip(parameters, gradients):\n        p_norm = unitwise_norm(params)\n        max_norm = tf.math.maximum(p_norm, eps) * clip_factor\n        grad_norm = unitwise_norm(grads)\n        clipped_grad = grads * (max_norm / tf.math.maximum(grad_norm, 1e-6))\n        new_grad = tf.where(grad_norm < max_norm, grads, clipped_grad)\n        new_grads.append(new_grad)\n    return new_grads","metadata":{"execution":{"iopub.status.busy":"2022-02-22T21:19:41.734947Z","iopub.execute_input":"2022-02-22T21:19:41.735397Z","iopub.status.idle":"2022-02-22T21:19:41.744013Z","shell.execute_reply.started":"2022-02-22T21:19:41.735361Z","shell.execute_reply":"2022-02-22T21:19:41.743494Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"class SupConLoss(tf.keras.losses.Loss):\n    def __init__(self, tau=1.0, name=None):\n        super().__init__(name=name)\n        self.tau = tau\n        \n    def __call__(self, y, z, sample_weight=None):\n        z, _ = tf.linalg.normalize(z, ord=2, axis=-1)\n        u = z @ tf.transpose(z) / self.tau\n        v = y @ tf.transpose(y)\n        v = v / tf.math.reduce_sum(v, axis=-1, keepdims=True)\n        return tf.nn.softmax_cross_entropy_with_logits(v, u)","metadata":{"execution":{"iopub.status.busy":"2022-02-22T21:19:41.744962Z","iopub.execute_input":"2022-02-22T21:19:41.745520Z","iopub.status.idle":"2022-02-22T21:19:41.757544Z","shell.execute_reply.started":"2022-02-22T21:19:41.745490Z","shell.execute_reply":"2022-02-22T21:19:41.756756Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"from tqdm.notebook import tqdm\n\nbsize = 1024\npaths = list(tqdm(data_root.joinpath(\"images\").glob(\"**/*.jpg\"), total=105100))\ntypes = [\n    label_names.index(articles[int(path.stem)][\"product_type_name\"])\n    for path in paths\n]\npaths = list(map(str, paths))\ntypes = tf.lookup.StaticHashTable(\n    tf.lookup.KeyValueTensorInitializer(\n        tf.constant(paths), tf.constant(types)\n    ),\n    default_value=-1\n)\nsteps = int(len(paths) / bsize + 0.5)\n\ndef deserialize(path):\n    x = tf.io.decode_jpeg(tf.io.read_file(path), channels=3)\n    x = tf.image.resize(x, (64, 64))\n    y = types.lookup(path)\n    return x, tf.one_hot(y, len(label_names))\n\ndef dataset(paths, training=True):\n    ds = (\n        tf.data.Dataset.from_tensor_slices(paths)\n        .map(deserialize)\n        .prefetch(tf.data.AUTOTUNE)\n    )\n    if training:\n        ds = ds.shuffle(bsize).cache()\n        ds = ds.batch(bsize)\n    return ds\n\ntrain_ds = dataset(paths[2048:])\ntest_ds = dataset(paths[:2048], training=False)","metadata":{"execution":{"iopub.status.busy":"2022-02-22T21:19:41.759058Z","iopub.execute_input":"2022-02-22T21:19:41.759543Z","iopub.status.idle":"2022-02-22T21:20:08.179010Z","shell.execute_reply.started":"2022-02-22T21:19:41.759504Z","shell.execute_reply":"2022-02-22T21:20:08.178274Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"embed_images, embed_labels = map(tf.stack, zip(*test_ds))\nembed_labels = [label_names[tf.argmax(y)] for y in embed_labels]\nembed_images.shape","metadata":{"execution":{"iopub.status.busy":"2022-02-22T21:20:08.180340Z","iopub.execute_input":"2022-02-22T21:20:08.180569Z","iopub.status.idle":"2022-02-22T21:20:29.918017Z","shell.execute_reply.started":"2022-02-22T21:20:08.180542Z","shell.execute_reply":"2022-02-22T21:20:29.916604Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"from comet_ml import Experiment\nfrom kaggle_secrets import UserSecretsClient\nuser_secrets = UserSecretsClient()\n\nmonitor = Experiment(\n    api_key=user_secrets.get_secret(\"comet-api-key\"),\n    project_name=\"hmrec\",\n    workspace=\"kavorite\",\n)\n\nclip_factor = 0.01\ntemperature = 0.05\nlatent_dims = 128\nmax_lr = 0.02 * bsize / 256\nmin_lr = max_lr * 1e-3\nepochs = 24\nmonitor.log_parameters({\n    \"temperature\": temperature,\n    \"latent_dims\": latent_dims,\n    \"max_lr\": max_lr,\n    \"min_lr\": min_lr,\n    \"clip_factor\": clip_factor,\n    \"epochs\": epochs,\n    \"batch_size\": bsize,\n})\n\n_, embed_image_url = monitor.create_embedding_image(\n    embed_images, image_size=tf.shape(embed_images)[-3:].numpy().tolist()\n)\nsup_con = SupConLoss(temperature)\nencoder = build_encoder(latent_dims)\ndecoder = build_decoder(latent_dims)\nopt_config = (steps * epochs, min_lr, max_lr)\nencoder_opt = build_optimizer(*opt_config)\ndecoder_opt = build_optimizer(*opt_config)\nacc1 = tf.keras.metrics.CategoricalAccuracy()\nacc5 = tf.keras.metrics.TopKCategoricalAccuracy(k=5)","metadata":{"execution":{"iopub.status.busy":"2022-02-22T21:20:29.919230Z","iopub.execute_input":"2022-02-22T21:20:29.919454Z","iopub.status.idle":"2022-02-22T21:20:41.599421Z","shell.execute_reply.started":"2022-02-22T21:20:29.919416Z","shell.execute_reply":"2022-02-22T21:20:41.598483Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"import time\n\ndef current_lr():\n    sched = encoder_opt.lr\n    steps = encoder_opt.iterations\n    return sched(steps) if callable(sched) else sched\n\n\nlast_viz = 0.0\nwith tqdm(total=steps * epochs) as progress:\n    for x, y in train_ds.repeat(epochs):\n        with tf.GradientTape() as tape:\n            z = encoder(x, training=True)\n            encoder_loss = tf.reduce_mean(sup_con(y, z))\n            \n        grads = tape.gradient(encoder_loss, encoder.trainable_weights)\n        clipped = adaptive_clip_grad(encoder.trainable_weights, grads, clip_factor)\n        encoder_opt.apply_gradients(\n            (g, w) if w.name.startswith(\"projection_head/\") else (c, w) \n            for g, c, w in zip(grads, clipped, encoder.trainable_weights)\n        )\n        \n        with tf.GradientTape() as tape:\n            y_hat = decoder(z)\n            decoder_loss = tf.reduce_mean(\n                tf.nn.softmax_cross_entropy_with_logits(y, y_hat / temperature)\n            )\n        \n        acc1.update_state(y, y_hat)\n        acc5.update_state(y, y_hat)\n        progress.set_description(\n            f\"loss = {encoder_loss.numpy():.3g}\"\n            + f\" acc@1 = {acc1.result():.3g}\"\n            + f\" acc@5 = {acc5.result():.3g}\"\n        )\n        progress.update()\n        # TODO: slow-cook encoder backbone during final epochs\n        epoch = tf.math.floor(encoder_opt.iterations / steps).numpy().astype(int)\n        grads = tape.gradient(decoder_loss, decoder.trainable_weights)\n        decoder_opt.apply_gradients(zip(grads, decoder.trainable_weights))\n        if (encoder_opt.iterations or -1) % steps == 0:\n            monitor.log_epoch_end(epochs)\n        monitor.log_metrics({\n            \"decoder_acc1\": acc1.result(),\n            \"decoder_acc5\": acc5.result(),\n            \"encoder_loss\": encoder_loss,\n            \"learning_rate\": current_lr(),\n        }, step=encoder_opt.iterations)\n        now = time.time()\n        if now - last_viz > 300:\n            last_viz = now\n            monitor.log_embedding(\n                vectors=encoder(embed_images, training=False).numpy().tolist(),\n                labels=embed_labels,\n                image_data=embed_image_url,\n                image_size=tf.shape(embed_images)[-3:].numpy().tolist(),\n                title=f\"Step {encoder_opt.iterations.numpy().astype(int)} Encodings\"\n            )","metadata":{"execution":{"iopub.status.busy":"2022-02-22T21:20:41.600877Z","iopub.execute_input":"2022-02-22T21:20:41.601120Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"import matplotlib.pyplot as plt\nimport matplotlib.cm as cm\nimport tensorflow_probability as tfp\n\ndef pca(mat, k=2):\n    loc, var = tf.nn.moments(mat, -1, keepdims=True)\n    mat = (mat - loc) / var\n    cov = tfp.stats.covariance(mat)\n    eigvals, eigvecs = tf.linalg.eigh(cov)\n    eigvecs = tf.gather(eigvecs, tf.argsort(eigvals)[::-1])\n    return mat @ eigvecs[:, :k]\n\nlabel = tf.argmax(y, -1)\norder = tf.argsort(label)\nlabel = tf.gather(label, order)\npreds = pca(tf.gather(z, order))\n_, _, group = tf.unique_with_counts(label)\ncolor = cm.nipy_spectral(tf.linspace(0, 1, len(group)))\nplt.figure(figsize=(12, 8))\nplt.title(\"Predictions\")\nfor lname, color, vectors in zip(label_names, color, tf.split(preds, group)):\n    if len(vectors) > 1:\n        plt.scatter(vectors[0], vectors[1], color=color, label=lname)\nplt.legend()","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"u = tf.nn.softmax(z @ tf.transpose(z))\nv = y @ tf.transpose(y) / tf.reduce_sum(y, axis=-1)\nplt.imshow(v - u)","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"encoder.save(\"encoder.h5\", include_optimizer=False)","metadata":{"trusted":true},"execution_count":null,"outputs":[]}]}