{"cells":[{"metadata":{},"cell_type":"markdown","source":"https://github.com/ayulockin/SwAV-TF  \nhttps://github.com/ayulockin/SwAV-TF/blob/master/Train_SwAV_40_epochs.ipynb"},{"metadata":{},"cell_type":"markdown","source":"# SwAV-TF\n- 本家の実装とは少し違う\n\n    - キューを使用しない。本家の実装ではプロトタイプの数がバッチサイズよりも大きくなったときに小さなバッチを使用する場合、本家はキューを維持\n    - 224x224の解像度の2つのクロップと、96x96の解像度の3つのクロップを使用。マルチクロップの提案された設定とは違う\n    - 15個のプロトタイプを使用。元の論文では、著者はImageNetデータセットに3000のプロトタイプを使用 \n    - 基本学習率0.1のコサイン減衰スケジュールとともにSGDを使用。本家の実装は習率スケジュールにウォームアップとコサイン減衰の組み合わせ\n"},{"metadata":{},"cell_type":"markdown","source":"\nSwAVは、同じ画像からのクロップ画像たちが属するクラスター(プロトタイプベクトル)は同じ、というのを学習\n\nこのクラスターのラベル(ソフトラベル)はSinkhorn-Knoppアルゴリズムと呼ばれるアルゴリズムで各バッチで毎回生成\n\nhttps://qiita.com/omiita/items/a7429ec42e4eef4b6a4d"},{"metadata":{},"cell_type":"markdown","source":"https://wandb.ai/authors/swav-tf/reports/Unsupervised-Visual-Representation-Learning-with-SwAV--VmlldzoyMjg3Mzg\n\nSwAVのアーキテクチャ\n\n![swav](https://i.ibb.co/TtSW4Fd/figure-3.png)\n\n![swav2](https://i.ibb.co/2FGDvd6/figure-6.png)\n\n![swav3](https://i.ibb.co/jgm7J81/figure-7.png)\n\n- 同じ画像をランダムにトリミングして高解像度（例：224x224）低解像度（例：96x96）のマルチクロップ画像にcolor distortion, random flipping, and random grayscalingなどのaugmentation を順番に適用（SimCLR のaugmentation）\n- CNN（ResNet50）で埋め込み（最後のグローバル平均プーリングレイヤーからの出力）ベクトルを取得\n- この埋め込みベクトルを浅い非線形ネットワークに送り、その出力が射影ベクトルZ\n- 射影ベクトルZを単一の線形層に渡す。つまり、このレイヤーには非線形性が含まれていません。レイヤーの出力は、Zとプロトタイプの間の内積。このレイヤーの関連する「重み」マトリックス（バックプロパゲーション中に更新されたもの）は、学習可能なプロトタイプバンクと見なすことができる\n- Sinkhorn Knopp アルゴリズムを使用して同じ画像の2つの別々のビュー間でスワップされた予測問題を設定"},{"metadata":{"_uuid":"d629ff2d2480ee46fbb7e2d37f6b5fab8052498a","_cell_guid":"79c7e3d0-c299-4dcb-8224-4455121ee9b0","trusted":true},"cell_type":"code","source":"#!git clone https://github.com/ayulockin/SwAV-TF.git","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"# https://github.com/ayulockin/SwAV-TF/tree/master/utils\n# multicrop_dataset.py\nimport tensorflow as tf\nimport random\n\nAUTO = tf.data.experimental.AUTOTUNE\n\n# Reference: https://github.com/google-research/simclr/blob/master/data_util.py\n\n@tf.function\ndef gaussian_blur(image, kernel_size=23, padding='SAME'):\n    sigma = tf.random.uniform((1,))* 1.9 + 0.1\n\n    radius = tf.cast(kernel_size / 2, tf.int32)\n    kernel_size = radius * 2 + 1\n    x = tf.cast(tf.range(-radius, radius + 1), tf.float32)\n    blur_filter = tf.exp(\n        -tf.pow(x, 2.0) / (2.0 * tf.pow(tf.cast(sigma, tf.float32), 2.0)))\n    blur_filter /= tf.reduce_sum(blur_filter)\n    # One vertical and one horizontal filter.\n    blur_v = tf.reshape(blur_filter, [kernel_size, 1, 1, 1])\n    blur_h = tf.reshape(blur_filter, [1, kernel_size, 1, 1])\n    num_channels = tf.shape(image)[-1]\n    blur_h = tf.tile(blur_h, [1, 1, num_channels, 1])\n    blur_v = tf.tile(blur_v, [1, 1, num_channels, 1])\n    expand_batch_dim = image.shape.ndims == 3\n    if expand_batch_dim:\n        image = tf.expand_dims(image, axis=0)\n    blurred = tf.nn.depthwise_conv2d(\n        image, blur_h, strides=[1, 1, 1, 1], padding=padding)\n    blurred = tf.nn.depthwise_conv2d(\n        blurred, blur_v, strides=[1, 1, 1, 1], padding=padding)\n    if expand_batch_dim:\n        blurred = tf.squeeze(blurred, axis=0)\n    return blurred\n\n@tf.function\ndef color_jitter(x, s=0.5):\n    x = tf.image.random_brightness(x, max_delta=0.8*s)\n    x = tf.image.random_contrast(x, lower=1-0.8*s, upper=1+0.8*s)\n    x = tf.image.random_saturation(x, lower=1-0.8*s, upper=1+0.8*s)\n    x = tf.image.random_hue(x, max_delta=0.2*s)\n    x = tf.clip_by_value(x, 0, 1)\n    return x\n\n@tf.function\ndef color_drop(x):\n    x = tf.image.rgb_to_grayscale(x)\n    x = tf.tile(x, [1, 1, 3])\n    return x\n\n@tf.function\ndef custom_augment(image):\n    # Random flips\n    image = random_apply(tf.image.flip_left_right, image, p=0.5)\n    # Randomly apply gausian blur\n    image = random_apply(gaussian_blur, image, p=0.5)\n    # Randomly apply transformation (color distortions) with probability p.\n    image = random_apply(color_jitter, image, p=0.8)\n    # Randomly apply grayscale\n    image = random_apply(color_drop, image, p=0.2)\n\n    return image\n\n@tf.function\ndef random_resize_crop(image, min_scale, max_scale, crop_size):\n    # Conditional resizing\n    if crop_size == 224:\n        image_shape = 260\n        image = tf.image.resize(image, (image_shape, image_shape))\n    else:\n        image_shape = 160\n        image = tf.image.resize(image, (image_shape, image_shape))\n    # Get the crop size for given min and max scale\n    size = tf.random.uniform(shape=(1,), minval=min_scale*image_shape,\n        maxval=max_scale*image_shape, dtype=tf.float32)\n    size = tf.cast(size, tf.int32)[0]\n    # Get the crop from the image\n    crop = tf.image.random_crop(image, (size,size,3))\n    crop_resize = tf.image.resize(crop, (crop_size, crop_size))\n\n    return crop_resize\n\n@tf.function\ndef random_apply(func, x, p):\n    return tf.cond(\n        tf.less(tf.random.uniform([], minval=0, maxval=1, dtype=tf.float32),\n                tf.cast(p, tf.float32)),\n        lambda: func(x),\n        lambda: x)\n\n@tf.function\ndef scale_image(image):\n    image = tf.image.convert_image_dtype(image, tf.float32)\n    return image\n\n@tf.function\ndef tie_together(image, min_scale, max_scale, crop_size):\n    # Retrieve the image features\n    image = image['image']\n    # Scale the pixel values\n    image = scale_image(image)\n    # Random resized crops\n    image = random_resize_crop(image, min_scale,\n        max_scale, crop_size)\n    # Color distortions & Gaussian blur\n    image = custom_augment(image)\n\n    return image\n\ndef get_multires_dataset(dataset,\n    size_crops,\n    num_crops,\n    min_scale,\n    max_scale,\n    options=None):\n    loaders = tuple()\n    for i, num_crop in enumerate(num_crops):\n        for _ in range(num_crop):\n            loader = (\n                    dataset\n                    .shuffle(1024)\n                    .map(lambda x: tie_together(x, min_scale[i],\n                        max_scale[i], size_crops[i]), num_parallel_calls=AUTO)\n                )\n            if options!=None:\n                loader = loader.with_options(options)\n            loaders += (loader, )\n\n    return loaders\n\ndef shuffle_zipped_output(a,b,c,d,e):\n    listify = [a,b,c,d,e]\n    random.shuffle(listify)\n\n    return listify[0], listify[1], listify[2], \\\n        listify[3], listify[4]","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"# https://github.com/ayulockin/SwAV-TF/tree/master/utils\n# architecture.py\nfrom tensorflow.keras import layers\nimport tensorflow as tf\n\ndef get_resnet_backbone():\n    base_model = tf.keras.applications.ResNet50(\n        include_top=False, weights=None, input_shape=(None, None, 3)\n    )\n    base_model.trainabe = True\n\n    inputs = layers.Input((None, None, 3))\n    h = base_model(inputs, training=True)\n    h = layers.GlobalAveragePooling2D()(h)\n    backbone = tf.keras.models.Model(inputs, h)\n\n    return backbone\n\ndef get_projection_prototype(dense_1=1024, dense_2=96, prototype_dimension=10):\n    inputs = layers.Input((2048, ))\n    projection_1 = layers.Dense(dense_1)(inputs)\n    projection_1 = layers.BatchNormalization()(projection_1)\n    projection_1 = layers.Activation(\"relu\")(projection_1)\n\n    projection_2 = layers.Dense(dense_2)(projection_1)\n    projection_2_normalize = tf.math.l2_normalize(projection_2, axis=1, name='projection')\n\n    prototype = layers.Dense(prototype_dimension, use_bias=False, name='prototype')(projection_2_normalize)\n\n    return tf.keras.models.Model(inputs=inputs,\n        outputs=[projection_2_normalize, prototype])","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"import tensorflow as tf\nimport tensorflow_datasets as tfds\n\nimport matplotlib.pyplot as plt\nimport numpy as np\nimport random\nimport time\nimport os\n\nfrom itertools import groupby\nfrom tqdm import tqdm\n\ntf.random.set_seed(666)\nnp.random.seed(666)\n\ntfds.disable_progress_bar()\n\nprint(\"Tensorflow version \" + tf.__version__)","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# params"},{"metadata":{"trusted":true},"cell_type":"code","source":"#epochs = 40\nepochs = 25  # 9時間以内に収めるため。1epoch~20分\n\nBATCH_SIZE = 32\n\nn_classes = 5  # キャッサバデータは5クラス\nn_prototype = n_classes * (3000 / 1000)  # プロトタイプの数.論文の比率と合わせる\n\n# エポック数減らして実行テスト\n#DEBUG = True\nDEBUG = False\nif DEBUG:\n    epochs = 2\n    print(\"DEBUG\")","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# Cassava data\n- https://www.kaggle.com/jessemostipak/getting-started-tpus-cassava-leaf-disease"},{"metadata":{"trusted":true},"cell_type":"code","source":"def decode_image(image):\n    image = tf.image.decode_jpeg(image, channels=3)\n    return image","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"def read_tfrecord(example, labeled):\n    tfrecord_format = {\n        \"image\": tf.io.FixedLenFeature([], tf.string),\n        \"target\": tf.io.FixedLenFeature([], tf.int64)\n    } if labeled else {\n        \"image\": tf.io.FixedLenFeature([], tf.string),\n        \"image_name\": tf.io.FixedLenFeature([], tf.string)\n    }\n    example = tf.io.parse_single_example(example, tfrecord_format)\n    image = decode_image(example['image'])\n    if labeled:\n        label = tf.cast(example['target'], tf.int32)\n        return {\"image\": image, \"label\": label}\n    idnum = example['image_name']\n    return {\"image\": image, \"label\": idnum}","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"from functools import partial\n\nAUTOTUNE = tf.data.experimental.AUTOTUNE\n\ndef load_dataset(filenames, labeled=True, ordered=False):\n    ignore_order = tf.data.Options()\n    if not ordered:\n        ignore_order.experimental_deterministic = False # disable order, increase speed\n    dataset = tf.data.TFRecordDataset(filenames, num_parallel_reads=AUTOTUNE) # automatically interleaves reads from multiple files\n    dataset = dataset.with_options(ignore_order) # uses data as soon as it streams in, rather than in its original order\n    dataset = dataset.map(partial(read_tfrecord, labeled=labeled), num_parallel_calls=AUTOTUNE)\n    return dataset","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"import re\n\ndef count_data_items(filenames):\n    n = [int(re.compile(r\"-([0-9]*)\\.\").search(filename).group(1)) for filename in filenames]\n    return np.sum(n)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"from kaggle_datasets import KaggleDatasets\nfrom sklearn.model_selection import train_test_split\n\nGCS_PATH = KaggleDatasets().get_gcs_path()\n\ntrain_tfrecs = tf.io.gfile.glob(GCS_PATH + '/train_tfrecords/ld_train*.tfrec')\n    \nTRAINING_FILENAMES, _ = train_test_split(train_tfrecs,\n                                         test_size=0.35, \n                                         random_state=5)\n\nNUM_TRAINING_IMAGES = count_data_items(TRAINING_FILENAMES)\n\nprint(f\"NUM_TRAINING_IMAGES: {NUM_TRAINING_IMAGES}\")","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# Multi Crop Resize Data Augmentation"},{"metadata":{"trusted":true},"cell_type":"code","source":"# Configs\nSIZE_CROPS = [224, 96]\nNUM_CROPS = [2, 3]\nMIN_SCALE = [0.14, 0.05] \nMAX_SCALE = [1., 0.14]\n\n# Experimental options\noptions = tf.data.Options()\noptions.experimental_optimization.noop_elimination = True\noptions.experimental_optimization.map_vectorization.enabled = True\noptions.experimental_optimization.apply_default_optimizations = True\noptions.experimental_deterministic = False\noptions.experimental_threading.max_intra_op_parallelism = 1","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"dataset = load_dataset(TRAINING_FILENAMES, labeled=True) \n\ntrainloaders = get_multires_dataset(dataset,\n    size_crops=SIZE_CROPS,\n    num_crops=NUM_CROPS,\n    min_scale=MIN_SCALE,\n    max_scale=MAX_SCALE,\n    options=options)\n\ntrainloaders","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"# Prepare the final data loader\n\nAUTO = tf.data.experimental.AUTOTUNE\n\n# Zipping \ntrainloaders_zipped = tf.data.Dataset.zip(trainloaders)\n\n# Final trainloader\ntrainloaders_zipped = (\n    trainloaders_zipped\n    .batch(BATCH_SIZE)\n    .prefetch(AUTO)\n)\n\nim1, im2, im3, im4, im5 = next(iter(trainloaders_zipped))\nprint(im1.shape, im2.shape, im3.shape, im4.shape, im5.shape)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"plt.figure(figsize=(15, 15))\nfor i, image_batch  in enumerate([im1, im2, im3, im4, im5]):\n    ax = plt.subplot(4, 5, i + 1)\n    plt.imshow(image_batch[0])\n    ax = plt.subplot(4, 5, 5 + i + 1)\n    plt.imshow(image_batch[1])\n    ax = plt.subplot(4, 5, 10 + i + 1)\n    plt.imshow(image_batch[2])\n    ax = plt.subplot(4, 5, 15 + i + 1)\n    plt.imshow(image_batch[3])","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# Model Architecture"},{"metadata":{"trusted":true},"cell_type":"code","source":"feature_backbone = get_resnet_backbone()\nfeature_backbone.summary()","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"projection_prototype = get_projection_prototype(n_prototype)\nprojection_prototype.summary()","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"embedding_batch = feature_backbone(im1)\nembedding_batch.shape","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"projection, prototype = projection_prototype(embedding_batch)\nprojection.shape, prototype.shape","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# Sinkhorn Knopp for Cluster Assignment\nReference: A.1 from https://arxiv.org/abs/2006.09882"},{"metadata":{"trusted":true},"cell_type":"code","source":"def sinkhorn(sample_prototype_batch, n_iters=3):\n    Q = tf.transpose(tf.exp(sample_prototype_batch/0.05))\n    Q /= tf.keras.backend.sum(Q)\n    K, B = Q.shape\n\n    u = tf.zeros_like(K, dtype=tf.float32)\n    r = tf.ones_like(K, dtype=tf.float32) / K\n    c = tf.ones_like(B, dtype=tf.float32) / B\n\n    for _ in range(n_iters):\n        u = tf.keras.backend.sum(Q, axis=1)\n        Q *= tf.expand_dims((r / u), axis=1)\n        Q *= tf.expand_dims(c / tf.keras.backend.sum(Q, axis=0), 0)\n\n    final_quantity = Q / tf.keras.backend.sum(Q, axis=0, keepdims=True)\n    final_quantity = tf.transpose(final_quantity)\n\n    return final_quantity","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"# Check\nfinal_q = sinkhorn(prototype)\nfinal_q.shape","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# Train Step"},{"metadata":{"trusted":true},"cell_type":"code","source":"# @tf.function\n# Reference: https://github.com/facebookresearch/swav/blob/master/main_swav.py\ndef train_step(input_views, feature_backbone, projection_prototype, \n               optimizer, crops_for_assign, temperature):\n    # ============ retrieve input data ... ============\n    im1, im2, im3, im4, im5 = input_views \n    inputs = [im1, im2, im3, im4, im5]\n    batch_size = inputs[0].shape[0]\n\n    # ============ create crop entries with same shape ... ============\n    crop_sizes = [inp.shape[1] for inp in inputs] # list of crop size of views\n    unique_consecutive_count = [len([elem for elem in g]) for _, g in groupby(crop_sizes)] # equivalent to torch.unique_consecutive\n    idx_crops = tf.cumsum(unique_consecutive_count)\n    \n    # ============ multi-res forward passes ... ============\n    start_idx = 0\n    with tf.GradientTape() as tape:\n        for end_idx in idx_crops:\n            concat_input = tf.stop_gradient(tf.concat(inputs[start_idx:end_idx], axis=0))\n            _embedding = feature_backbone(concat_input) # get embedding of same dim views together\n            if start_idx == 0:\n                embeddings = _embedding # for first iter\n            else:\n                embeddings = tf.concat((embeddings, _embedding), axis=0) # concat all the embeddings from all the views\n            start_idx = end_idx\n        \n        projection, prototype = projection_prototype(embeddings) # get normalized projection and prototype\n        projection = tf.stop_gradient(projection)\n\n        # ============ swav loss ... ============\n        # https://github.com/facebookresearch/swav/issues/19\n        loss = 0\n        for i, crop_id in enumerate(crops_for_assign): # crops_for_assign = [0,1]\n            with tape.stop_recording():\n                out = prototype[batch_size * crop_id: batch_size * (crop_id + 1)]\n                \n                # get assignments\n                q = sinkhorn(out) # sinkhorn is used for cluster assignment\n            \n            # cluster assignment prediction\n            subloss = 0\n            for v in np.delete(np.arange(np.sum(NUM_CROPS)), crop_id): # (for rest of the portions compute p and take cross entropy with q)\n                p = tf.nn.softmax(prototype[batch_size * v: batch_size * (v + 1)] / temperature) \n                subloss -= tf.math.reduce_mean(tf.math.reduce_sum(q * tf.math.log(p), axis=1))\n            loss += subloss / tf.cast((tf.reduce_sum(NUM_CROPS) - 1), tf.float32)\n        \n        loss /= len(crops_for_assign)\n\n    # ============ backprop ... ============\n    variables = feature_backbone.trainable_variables + projection_prototype.trainable_variables\n    gradients = tape.gradient(loss, variables)\n    optimizer.apply_gradients(zip(gradients, variables))\n\n    return loss","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# Training Loop"},{"metadata":{"trusted":true},"cell_type":"code","source":"def train_swav(feature_backbone, \n               projection_prototype, \n               dataloader, \n               optimizer, \n               crops_for_assign,\n               temperature, \n               epochs=50):\n  \n    step_wise_loss = []\n    epoch_wise_loss = []\n    \n    for epoch in tqdm(range(epochs)):\n        w = projection_prototype.get_layer('prototype').get_weights()\n        w = tf.transpose(w)\n        w = tf.math.l2_normalize(w, axis=1)\n        projection_prototype.get_layer('prototype').set_weights(tf.transpose(w))\n\n        for i, inputs in enumerate(dataloader):\n            loss = train_step(inputs, feature_backbone, projection_prototype, \n                              optimizer, crops_for_assign, temperature)\n            step_wise_loss.append(loss)\n\n        epoch_wise_loss.append(np.mean(step_wise_loss))\n        #wandb.log({'epoch': epoch, 'loss':np.mean(step_wise_loss)})\n        \n        if epoch % 5 == 0:\n            print(\"epoch: {} loss: {:.3f}\".format(epoch + 1, np.mean(step_wise_loss)))\n\n    return epoch_wise_loss, [feature_backbone, projection_prototype]","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"%%time\n# ============ initialize the networks and the optimizer ... ============\nfeature_backbone = get_resnet_backbone()\nprojection_prototype = get_projection_prototype(10)\n\nlr_decayed_fn = tf.keras.optimizers.schedules.PolynomialDecay(\n    initial_learning_rate=0.1,\n    decay_steps=500,\n    end_learning_rate=0.001,\n    power=0.5)\nopt = tf.keras.optimizers.SGD(learning_rate=lr_decayed_fn)\n\n# ================= initialize wandb ======================\n#wandb.init(entity='authors', project='swav-tf', id='40-epochs')\n\n# ============ train for 40 epochs ... ============\nepoch_wise_loss, models = train_swav(feature_backbone, \n    projection_prototype, \n    trainloaders_zipped, \n    opt, \n    crops_for_assign=[0, 1],\n    temperature=0.1, \n    epochs=epochs\n)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"plt.plot(epoch_wise_loss)\nplt.show()","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"# Serialize the models\nfeature_backbone, projection_prototype = models\nfeature_backbone.save_weights('feature_backbone.h5')\nprojection_prototype.save_weights('projection_prototype.h5')","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}