{"cells":[{"metadata":{},"cell_type":"markdown","source":"## 参考\nhttps://www.kaggle.com/abhinand05/vision-transformer-vit-tutorial-baseline"},{"metadata":{"_kg_hide-input":true,"_kg_hide-output":true,"trusted":true},"cell_type":"code","source":"!cp -r ../input/vittutorialillustrations/* ./ \n\n!pip install nb_black\n%load_ext nb_black","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# Introduction\n\nこのノートは、タイトルからもわかるように、基本的には2つのパートに分かれています。\n\n<a href=\"#Vision-Transformers:-A-gentle-introduction\">1. Vision Transformer: 優しい紹介</a> <br>\n<a href=\"#Vision-Transformer-Implementation-in-PyTorch\">2. PyTorchでの実装</a>\n\n**今回のコンテストのために、PyTorch での ViT の実装に入る前に、Vision Transformers の基本的な考え方とその仕組みについて簡単に説明します。**\n\nコードだけに興味がある方は、このノートの第二章を読み飛ばしても構いません。 実装は、[rwightman/pytorch-image-models](https://github.com/rwightman/pytorch-image-models) ライブラリのお陰で大きく変わることはありませんが、このライブラリには、事前学習された重みを含むすべてのモデルの実装が含まれています。"},{"metadata":{},"cell_type":"markdown","source":"# <font size=4 color='blue'>このノートブックが便利だと思ったら、私はもっとそのようなノートブックを書くためにやる気にさせるUpvoteを残してください。</font>"},{"metadata":{},"cell_type":"markdown","source":"# Vision Transformers: A gentle introduction\n\nVision Transformersは、2020年10月下旬にGoogle Brainチームが発表した論文 [AN IMAGE IS WORTH 16X16 WORDS:\nTRANSFORMERS FOR IMAGE RECOGNITION AT SCALE](https://arxiv.org/pdf/2010.11929.pdf) で初めて紹介されました。\n\nViTの仕組みを理解するためには、当然ながらtransformerがどのように機能し、どのような問題を解決したのかについての予備知識が必要です。ここでは、transformerがどのように機能するのかを簡単に紹介してから、目の前のトピックであるViTの詳細に入っていこうと思います。\n\n![ViT-Illustration](vision-transformer.png)\n\nもしあなたがNLP（自然言語処理）を初めて知り、トランスフォーマーモデルについてもっと知りたいと思っていて、実際にどのように機能するかについての公正な直観を得たいと思っているなら、 [Jay Allamar](https://jalammar.github.io/)の素晴らしいブログ記事をチェックしてみることをお勧めします。上の画像も彼のブログ記事からインスピレーションを得ています。"},{"metadata":{},"cell_type":"markdown","source":"## Transformers: A brief overview\n\n> **すでにTransformerを理解されている方は、ご自由に読み飛ばしてください。**\n\nトランスフォーマーモデルは、私たちが知っているように自然言語処理に革命をもたらしました。最初に導入されたとき、トランスフォーマーモデルは複数のNLP記録を更新し、当時のState of the Artを押し進めていました。今では、現代のNLPタスクのデファクトスタンダードとなっており、LSTMやGRUのような前世代のモデルと比較すると、驚くほどのパフォーマンス向上をもたらします。\n\nNLPの風景を一変させた最も重要な論文は、[\"Attention is all you need\"](https://arxiv.org/pdf/1706.03762.pdf) という論文です。この論文で紹介されたのが、トランスフォーマーアーキテクチャです。\n\n### **Motivations:**\n\n当時、系列やNLPタスクのための既存のモデルは、ほとんどがRNNを使用していました。 **これらのネットワークの問題点は、長期的な依存関係を捉えることができないことでした。** \n\nLSTMやGRU - RNNの亜種は依存関係をキャプチャすることができましたが、それにも限界がありました。 \n\nそこで、トランスフォーマーの背後にある主なインスピレーションは、この再帰を取り除き、ほぼすべての依存性をキャプチャすることでした。これは、self-attention（マルチヘッド）と呼ばれる注意メカニズムの変形を使用して達成されたもので、成功には非常に重要です。トランフォーマーモデルのもう一つの利点は、高度な並列化が可能であることです。  \n\n\n### Transformer Architecture\n**注: アーキテクチャ図には、説明中の対応するステップが注記されています。**\n\n![TranformerArchitecture](transformer-arch.png)\n\n- Transformerには、上図の左側にあるデコーダと右側にあるエンコーダの2つの部分があります。 \n- ここでは機械翻訳をしていると想像してください。 \n- エンコーダは入力データ（文）を受け取り、入力の中間表現を生成します。 \n- デコーダはこの中間表現を段階的にデコードし、出力を生成します。しかし、違いはこれをどのように行っているかにあります。 \n- ViTでは、エンコーダのセクションを理解するだけで十分です。 \n\n> **注: ここでの説明は、アーキテクチャの背後にある直感についてのものです。より多くの数学的な詳細については、代わりにそれぞれの研究論文をチェックしてください。**\n\n### Tranformers: Step by step overview\n**(1)** 入力データは最初にベクトルに埋め込まれます。埋め込み層(embedding layer)は、各単語のために学習されたベクトル表現をつかむのに役立つ。\n\n**(2)** 次の段階では、位置エンコーディングが入力embeddingに注入される。これは、Tranformerが入力として渡されるシーケンスの順序（例えば文）を知らないからです\n\n**(3)** ここで、 multi-headed attentionが少し変わってきます。\n\n**Multi-headed-attention architecture:**\n![multi-headed-attn](multi-headed-attention.png)\n\n**(4)** Multi-Headed Attentionは3つの学習可能なベクトルで構成されています。Query, Key、Valueの3つのベクトルである。これは、検索（クエリ）すると、検索エンジンがクエリとキーを比較し、値で応答するという情報の再利用に由来すると言われています。\n\n**(5)** Q と K の表現は、ドット積行列の乗算を経て、ある単語が他のすべての単語にどれだけ注意を払わなければならないかを表すスコア行列を生成します。スコアが高ければ高いほど注目度が高く、逆もまた然りです。 \n\n**(6)** その後、スコア行列は、Q と K ベクトルの次元に応じてスケールダウンされる。これは，乗算が爆発的な効果をもたらす可能性があるため，より安定した勾配を確保するためです． \n\n(マスクの部分については、デコーダのセクションに到達したときに説明します)\n\n**(7)** 次に、注目度スコアを確率に変換するために、スコア行列をソフトマックス化します。明らかに、スコアが高いほど高くなり、低いほど低くなります。これにより、モデルがどの単語に注目すべきかを確実に判断できるようになります。 \n\n**(8)** 次に、確率を含む結果の行列に値ベクトルを乗算します。これにより、モデルが学習した確率スコアの高い単語がより重要になる。スコアの低い単語は効果的にかき消されて無関係になる。 \n\n**(9)** そして、QKベクトルとVベクトルの連結出力をLinear層に送り込み、さらに処理を行う。 \n\n**(10)** シーケンス内の各単語に対してSelf-Attentionが行われる。1つは他のものに依存しないので、Self-Attentionモジュールのコピーを使用して、これを**multi-headed**化して同時にすべてを処理することができます。 \n\n**(11)** その後、出力値ベクトルを連結し、入力層からの残差接続に加算し、その結果の再表現をLayernNormに渡して正規化する。(残差接続はネットワークを流れる勾配を助け、LayernNormは学習時間をわずかに短縮し、ネットワークを安定化させるのに役立ちます)\n\n**(12)** さらに、出力は、より豊かな表現を得るために、point-wise feed forward networkに渡されます。  \n\n**(13)**  出力は再びレイヤーノルム化され、前のレイヤーから残差が追加されます。 \n\n\n**注意: これでエンコーダのセクションは終わりですが、Vision Transformer を完全に理解するにはこれで十分です。デコーダ部分はエンコーディングレイヤーと非常に似ているので、理解するのはあなたにお任せします。**\n\n**(14)** エンコーダからの出力は、前の時間ステップ/ワードからの入力（もしあれば）と共にデコーダに送られ、出力はエンコーダからの出力と共に次のattention layerに送られる前に、マスクされたmulti headed attentionを受けます。 \n\n**(15)** Masked multi headed attention はリークがないことを確実にするために、デコード中にネットワークがシーケンス内で後から来る単語への可視性を持つべきではないために必要である。これは、スコアマトリクスの系列内で後から来る単語のエントリをマスクすることによって行われます。シーケンス内の現在の単語と前の単語は 1 で追加され、未来の単語のスコアは -inf で追加されます。これにより、確率を得るためにsoftmaxを実行する際に、系列内の将来の単語が0にかき消され、残りの単語は保持されます。 \n\n**(16)** ここにも残差接続があり、勾配の流れを改善しています。最後に、出力は線形層に送られ、確率の出力を得るためにソフトマックスされます。 "},{"metadata":{},"cell_type":"markdown","source":"## How Vision Tranformers works?\n\nTranformerの内部の働きを高いレベルでカバーしたところで、いよいよVision Tranformersに取り組む準備が整いました。 \n\nTranformerを画像に適用することは、以下の理由から常に困難なことでした。\n- 単語/文章/段落とは異なり、画像は基本的にはピクセルの形でより多くの情報を含んでいます。 \n- 現在のハードウェアでも、画像内のすべてのピクセルに注目することは非常に困難です。 \n- その代わりに、人気のある代替案は、局所的なattentionを利用することでした。 \n- 実際、CNNは畳み込みによって非常に似たようなことをしていて、モデルの層を深くしていくと受容野は本質的に大きくなりますが、Tranformerは「Tranformer」の性質上、常にCNNよりも計算量が多くなります。もちろん、CNNが現在のコンピュータビジョンの進歩にどれだけ貢献しているかは知っている。\n\nグーグルの研究者たちは論文の中で、コンピュータ・ビジョンの次の大きな一歩となりうる、これまでとは異なるものを提案している。彼らはCNNへの依存はもう必要ないかもしれないことを示しているのだ。では、Vision Tranformerについてもっと詳しく見ていこう。\n\n### Vision Transformer Architecture\n\n![vit-architecture](vit-arch.png)\n\n**(1)** トランスフォーマーのエンコーダ部分だけを使用していますが、画像をネットワークに送り込む方法に違いがあります。\n\n\n**(2)** 画像を固定サイズのパッチに分解しています。そのため、これらのパッチの1つは、論文で提案されているように、16x16または32x32の寸法にすることができます。パッチが多ければ多いほど、パッチ自体が小さくなるので、これらのネットワークを訓練するのがより簡単になります。したがって、私たちはタイトルにあるように、「画像は16x16ワードの価値がある」ということになります。 \n\n**(3)** その後、パッチは展開され（平坦化され）、ネットワークへの更なる処理のために送られます。\n\n**(4)** ここでのNNとは異なり、モデルはシーケンス内のサンプルの位置について何も考えていませんが、ここでは各サンプルは入力画像からのパッチです。 そのため、画像は**位置埋め込みベクトルと一緒に**エンコーダに送り込まれます。ここで注意しなければならないことは、位置埋め込みも学習可能なので、実際には位置に関係なくハードコードされたベクトルを送り込む必要はないということです。\n\n**(5)** BERTのように開始時に特別なトークンもあります。\n\n**(6)** 各画像パッチは、最初に大きなベクトルに展開（平坦化）され、学習可能な埋め込み行列と掛け合わされ、埋め込みパッチが作成されます。そして、これらの埋め込みパッチは位置埋め込みベクトルと結合され、それがトランスフォーマーに供給されます。 \n\n> **注意：ここから先はすべて標準的なトランスフォーマーと同じです。**\n\n**(7)** 唯一の違いは、デコーダの代わりにエンコーダからの出力が直接フィードフォワードニューラルネットワークに渡され、分類出力を得ることです。 \n\n### Things to note:\n- この論文は、ほとんどの場合、コンボリューションを完全に無視しています。 \n- しかし、彼らは、画像パッチのコンボリューション埋め込みを使用するViTのいくつかのバリエーションを使用しています。しかし、それは性能にあまり影響を与えていないようです。  \n- これを書いている時点では、Vision TransformersはImageNetの画像分類ベンチマークでトップになっています。 \n\n<img src=\"benchmarks-chart.png\" width=\"700\">\n<!-- ![BenchmarksChart](benchmarks-charpng) -->\n\n- この論文には他にも興味深いことがたくさんありますが、私にとって目立っていて、潜在的にCNNよりもトランスフォーマーの力を示しているのは、下の画像のように、レイヤーに対するattention distanceを示していることです。 \n\n\n<img src=\"attn-distance.png\" width=\"300\" height=\"300\">\n<br>\n\n- 上のグラフは、トランスフォーマーがネットワークの開始層から離れた領域にすでにattentionを払う能力を持っていることを示唆しています。"},{"metadata":{},"cell_type":"markdown","source":"# <font size=4 color='blue'>このノートブックが便利だと思ったら、私はもっとそのようなノートブックを書くためにやる気にさせるUpvoteを残してください。</font>"},{"metadata":{},"cell_type":"markdown","source":"## Vision Transformer Implementation in PyTorch\nビジョントランスフォーマーを理解したところで、  [this competition](https://www.kaggle.com/c/cassava-leaf-disease-classification)のベースラインモデルを構築してみましょう。\n\nまず、TPU と torch-image-models (timm) を使えるようにするために torch-xla をインストールします。"},{"metadata":{"_kg_hide-input":true,"_kg_hide-output":true,"trusted":true},"cell_type":"code","source":"!curl https://raw.githubusercontent.com/pytorch/xla/master/contrib/scripts/env-setup.py -o pytorch-xla-env-setup.py\n!python pytorch-xla-env-setup.py --version 1.7\n!pip install timm","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"import numpy as np\nimport pandas as pd\nimport matplotlib.pyplot as plt\n\nplt.style.use(\"ggplot\")\n\nimport torch\nimport torch.nn as nn\nimport torchvision.transforms as transforms\n\nimport torch_xla\nimport torch_xla.core.xla_model as xm\nimport torch_xla.distributed.xla_multiprocessing as xmp\nimport torch_xla.distributed.parallel_loader as pl\n\nimport timm\n\nimport gc\nimport os\nimport time\nimport random\nfrom datetime import datetime\n\nfrom PIL import Image\nfrom tqdm.notebook import tqdm\nfrom sklearn import model_selection, metrics","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"# For parallelization in TPUs\nos.environ[\"XLA_USE_BF16\"] = \"1\"\nos.environ[\"XLA_TENSOR_ALLOCATOR_MAXSIZE\"] = \"100000000\"","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"def seed_everything(seed):\n    \"\"\"\n    Seeds basic parameters for reproductibility of results\n    \n    Arguments:\n        seed {int} -- Number of the seed\n    \"\"\"\n    random.seed(seed)\n    os.environ[\"PYTHONHASHSEED\"] = str(seed)\n    np.random.seed(seed)\n    torch.manual_seed(seed)\n    torch.cuda.manual_seed(seed)\n    torch.backends.cudnn.deterministic = True\n    torch.backends.cudnn.benchmark = False\n\n\nseed_everything(1001)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"# general global variables\nDATA_PATH = \"../input/cassava-leaf-disease-classification\"\nTRAIN_PATH = \"../input/cassava-leaf-disease-classification/train_images/\"\nTEST_PATH = \"../input/cassava-leaf-disease-classification/test_images/\"\nMODEL_PATH = (\n    \"../input/vit-base-models-pretrained-pytorch/jx_vit_base_p16_224-80ecf9dd.pth\"\n)\n\n# model specific global variables\nIMG_SIZE = 224\nBATCH_SIZE = 16\nLR = 2e-05\nGAMMA = 0.7\nN_EPOCHS = 10","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"df = pd.read_csv(os.path.join(DATA_PATH, \"train.csv\"))\ndf.head()","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"df.info()","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"df.label.value_counts().plot(kind=\"bar\")","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"train_df, valid_df = model_selection.train_test_split(\n    df, test_size=0.1, random_state=42, stratify=df.label.values\n)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"class CassavaDataset(torch.utils.data.Dataset):\n    \"\"\"\n    Helper Class to create the pytorch dataset\n    \"\"\"\n\n    def __init__(self, df, data_path=DATA_PATH, mode=\"train\", transforms=None):\n        super().__init__()\n        self.df_data = df.values\n        self.data_path = data_path\n        self.transforms = transforms\n        self.mode = mode\n        self.data_dir = \"train_images\" if mode == \"train\" else \"test_images\"\n\n    def __len__(self):\n        return len(self.df_data)\n\n    def __getitem__(self, index):\n        img_name, label = self.df_data[index]\n        img_path = os.path.join(self.data_path, self.data_dir, img_name)\n        img = Image.open(img_path).convert(\"RGB\")\n\n        if self.transforms is not None:\n            image = self.transforms(img)\n\n        return image, label","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"# create image augmentations\ntransforms_train = transforms.Compose(\n    [\n        transforms.Resize((IMG_SIZE, IMG_SIZE)),\n        transforms.RandomHorizontalFlip(p=0.3),\n        transforms.RandomVerticalFlip(p=0.3),\n        transforms.RandomResizedCrop(IMG_SIZE),\n        transforms.ToTensor(),\n        transforms.Normalize((0.485, 0.456, 0.406), (0.229, 0.224, 0.225)),\n    ]\n)\n\ntransforms_valid = transforms.Compose(\n    [\n        transforms.Resize((IMG_SIZE, IMG_SIZE)),\n        transforms.ToTensor(),\n        transforms.Normalize((0.485, 0.456, 0.406), (0.229, 0.224, 0.225)),\n    ]\n)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"print(\"Available Vision Transformer Models: \")\ntimm.list_models(\"vit*\")","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"class ViTBase16(nn.Module):\n    def __init__(self, n_classes, pretrained=False):\n\n        super(ViTBase16, self).__init__()\n\n        self.model = timm.create_model(\"vit_base_patch16_224\", pretrained=False)\n        if pretrained:\n            self.model.load_state_dict(torch.load(MODEL_PATH))\n\n        self.model.head = nn.Linear(self.model.head.in_features, n_classes)\n\n    def forward(self, x):\n        x = self.model(x)\n        return x\n\n    def train_one_epoch(self, train_loader, criterion, optimizer, device):\n        # keep track of training loss\n        epoch_loss = 0.0\n        epoch_accuracy = 0.0\n\n        ###################\n        # train the model #\n        ###################\n        self.model.train()\n        for i, (data, target) in enumerate(train_loader):\n            # move tensors to GPU if CUDA is available\n            if device.type == \"cuda\":\n                data, target = data.cuda(), target.cuda()\n            elif device.type == \"xla\":\n                data = data.to(device, dtype=torch.float32)\n                target = target.to(device, dtype=torch.int64)\n\n            # clear the gradients of all optimized variables\n            optimizer.zero_grad()\n            # forward pass: compute predicted outputs by passing inputs to the model\n            output = self.forward(data)\n            # calculate the batch loss\n            loss = criterion(output, target)\n            # backward pass: compute gradient of the loss with respect to model parameters\n            loss.backward()\n            # Calculate Accuracy\n            accuracy = (output.argmax(dim=1) == target).float().mean()\n            # update training loss and accuracy\n            epoch_loss += loss\n            epoch_accuracy += accuracy\n\n            # perform a single optimization step (parameter update)\n            if device.type == \"xla\":\n                xm.optimizer_step(optimizer)\n\n                if i % 20 == 0:\n                    xm.master_print(f\"\\tBATCH {i+1}/{len(train_loader)} - LOSS: {loss}\")\n\n            else:\n                optimizer.step()\n\n        return epoch_loss / len(train_loader), epoch_accuracy / len(train_loader)\n\n    def validate_one_epoch(self, valid_loader, criterion, device):\n        # keep track of validation loss\n        valid_loss = 0.0\n        valid_accuracy = 0.0\n\n        ######################\n        # validate the model #\n        ######################\n        self.model.eval()\n        for data, target in valid_loader:\n            # move tensors to GPU if CUDA is available\n            if device.type == \"cuda\":\n                data, target = data.cuda(), target.cuda()\n            elif device.type == \"xla\":\n                data = data.to(device, dtype=torch.float32)\n                target = target.to(device, dtype=torch.int64)\n\n            with torch.no_grad():\n                # forward pass: compute predicted outputs by passing inputs to the model\n                output = self.model(data)\n                # calculate the batch loss\n                loss = criterion(output, target)\n                # Calculate Accuracy\n                accuracy = (output.argmax(dim=1) == target).float().mean()\n                # update average validation loss and accuracy\n                valid_loss += loss\n                valid_accuracy += accuracy\n\n        return valid_loss / len(valid_loader), valid_accuracy / len(valid_loader)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"def fit_tpu(\n    model, epochs, device, criterion, optimizer, train_loader, valid_loader=None\n):\n\n    valid_loss_min = np.Inf  # track change in validation loss\n\n    # keeping track of losses as it happen\n    train_losses = []\n    valid_losses = []\n    train_accs = []\n    valid_accs = []\n\n    for epoch in range(1, epochs + 1):\n        gc.collect()\n        para_train_loader = pl.ParallelLoader(train_loader, [device])\n\n        xm.master_print(f\"{'='*50}\")\n        xm.master_print(f\"EPOCH {epoch} - TRAINING...\")\n        train_loss, train_acc = model.train_one_epoch(\n            para_train_loader.per_device_loader(device), criterion, optimizer, device\n        )\n        xm.master_print(\n            f\"\\n\\t[TRAIN] EPOCH {epoch} - LOSS: {train_loss}, ACCURACY: {train_acc}\\n\"\n        )\n        train_losses.append(train_loss)\n        train_accs.append(train_acc)\n        gc.collect()\n\n        if valid_loader is not None:\n            gc.collect()\n            para_valid_loader = pl.ParallelLoader(valid_loader, [device])\n            xm.master_print(f\"EPOCH {epoch} - VALIDATING...\")\n            valid_loss, valid_acc = model.validate_one_epoch(\n                para_valid_loader.per_device_loader(device), criterion, device\n            )\n            xm.master_print(f\"\\t[VALID] LOSS: {valid_loss}, ACCURACY: {valid_acc}\\n\")\n            valid_losses.append(valid_loss)\n            valid_accs.append(valid_acc)\n            gc.collect()\n\n            # save model if validation loss has decreased\n            if valid_loss <= valid_loss_min and epoch != 1:\n                xm.master_print(\n                    \"Validation loss decreased ({:.4f} --> {:.4f}).  Saving model ...\".format(\n                        valid_loss_min, valid_loss\n                    )\n                )\n            #                 xm.save(model.state_dict(), 'best_model.pth')\n\n            valid_loss_min = valid_loss\n\n    return {\n        \"train_loss\": train_losses,\n        \"valid_losses\": valid_losses,\n        \"train_acc\": train_accs,\n        \"valid_acc\": valid_accs,\n    }","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"model = ViTBase16(n_classes=5, pretrained=True)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"def _run():\n    train_dataset = CassavaDataset(train_df, transforms=transforms_train)\n    valid_dataset = CassavaDataset(valid_df, transforms=transforms_valid)\n\n    train_sampler = torch.utils.data.distributed.DistributedSampler(\n        train_dataset,\n        num_replicas=xm.xrt_world_size(),\n        rank=xm.get_ordinal(),\n        shuffle=True,\n    )\n\n    valid_sampler = torch.utils.data.distributed.DistributedSampler(\n        valid_dataset,\n        num_replicas=xm.xrt_world_size(),\n        rank=xm.get_ordinal(),\n        shuffle=False,\n    )\n\n    train_loader = torch.utils.data.DataLoader(\n        dataset=train_dataset,\n        batch_size=BATCH_SIZE,\n        sampler=train_sampler,\n        drop_last=True,\n        num_workers=8,\n    )\n\n    valid_loader = torch.utils.data.DataLoader(\n        dataset=valid_dataset,\n        batch_size=BATCH_SIZE,\n        sampler=valid_sampler,\n        drop_last=True,\n        num_workers=8,\n    )\n\n    criterion = nn.CrossEntropyLoss()\n    #     device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n    device = xm.xla_device()\n    model.to(device)\n\n    lr = LR * xm.xrt_world_size()\n    optimizer = torch.optim.Adam(model.parameters(), lr=lr)\n\n    xm.master_print(f\"INITIALIZING TRAINING ON {xm.xrt_world_size()} TPU CORES\")\n    start_time = datetime.now()\n    xm.master_print(f\"Start Time: {start_time}\")\n\n    logs = fit_tpu(\n        model=model,\n        epochs=N_EPOCHS,\n        device=device,\n        criterion=criterion,\n        optimizer=optimizer,\n        train_loader=train_loader,\n        valid_loader=valid_loader,\n    )\n\n    xm.master_print(f\"Execution time: {datetime.now() - start_time}\")\n\n    xm.master_print(\"Saving Model\")\n    xm.save(\n        model.state_dict(), f'model_5e_{datetime.now().strftime(\"%Y%m%d-%H%M\")}.pth'\n    )","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"# Start training processes\ndef _mp_fn(rank, flags):\n    torch.set_default_tensor_type(\"torch.FloatTensor\")\n    a = _run()\n\n\n# _run()\nFLAGS = {}\nxmp.spawn(_mp_fn, args=(FLAGS,), nprocs=8, start_method=\"fork\")","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"## Thanks a lot for reading all the way\n\n# <font size=4 color='blue'>If you find this notebook useful, leave an upvote, that motivates me to write more such notebooks.</font>"}],"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}