{"cells":[{"metadata":{},"cell_type":"markdown","source":"Toturials:\n\nhttps://github.com/pytorch/xla/\n\nhttp://pytorch.org/xla/release/1.5/index.html"},{"metadata":{"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 nightly --apt-packages libomp5 libopenblas-dev","execution_count":null,"outputs":[]},{"metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","trusted":true},"cell_type":"code","source":"import os\nimport time\nimport numpy as np \nimport pandas as pd\nfrom collections import deque\nimport matplotlib.pyplot as plt\n\nimport warnings\nwarnings.filterwarnings(\"ignore\")\n\nfor dirname, _, filenames in os.walk('/kaggle/input'):\n    for filename in filenames:\n        print(os.path.join(dirname, filename))\n","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"!echo $TPU_NAME","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"import torch\nimport torch.nn as nn\nfrom torch.utils.data import Dataset, DataLoader\nfrom torch.utils.data.distributed import DistributedSampler\n\nimport torchvision\nimport torchvision.transforms as T \n\n# imports the torch_xla package\nimport torch_xla\nimport torch_xla.core.xla_model as xm\nimport torch_xla.distributed.parallel_loader as pl\nimport torch_xla.distributed.xla_multiprocessing as xmp","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# Build custom dataset object"},{"metadata":{"trusted":true},"cell_type":"code","source":"train_df = pd.read_csv(dirname+'/train.csv')\nlabels = train_df['label'].values\nimg_arrs = train_df.iloc[:,1:].values\n\nidx = 10\nimg_arr = img_arrs[idx]/255.\n\nplt.imshow(img_arr.reshape(28,28), cmap='gray')\nplt.show()","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"transforms = T.Compose([T.ToPILImage(), T.Resize((224,224)), T.ToTensor()])\nimg_tensor = transforms(np.float32(img_arr.reshape(28,28,1)))\nplt.imshow(img_tensor[0],cmap='gray')\nplt.show()","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"class mnistDataset(Dataset):\n    def __init__(self, img_arrs, labels = None):\n        self.img_arrs = img_arrs\n        self.labels = labels\n        self.transforms = T.Compose([T.ToPILImage(), T.Resize((224,224)), T.ToTensor()])\n        \n    def __getitem__(self, idx):\n        img_arr = self.img_arrs[idx]/255 \n        original_img = img_arr.reshape(28,28,1)\n        img_tensor = self.transforms(np.float32(original_img))\n        if self.labels is None:\n            return img_tensor\n        else:\n            target_tensor = torch.tensor(self.labels[idx], dtype=torch.long)\n            return img_tensor, target_tensor\n    \n    def __len__(self):\n        return len(self.img_arrs)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"dataset = mnistDataset(img_arrs,labels)\nprint(dataset.__len__())\nbatch_size = 32\ntrain_loader = DataLoader(dataset, batch_size = batch_size, shuffle = True, drop_last = True)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"trainiter = iter(train_loader)\nimg_tensor, target_tensor = next(trainiter)\nprint(img_tensor.size())\nprint(target_tensor.size())\n\nfig = plt.figure(figsize=(25, 8))\nplot_size = 32\nfor idx in np.arange(plot_size):\n    ax = fig.add_subplot(4, plot_size/4, idx+1, xticks=[], yticks=[])\n    ax.imshow(img_tensor[idx][0], cmap='gray')","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# Adjust AlexNet for MNIST data set"},{"metadata":{"trusted":true},"cell_type":"code","source":"net = torchvision.models.alexnet(num_classes=10)\nlayers = [nn.Conv2d(1, 64, kernel_size = 11, stride = 4, padding = 2)] + list(net.features.children())[1:]\nnet.features = nn.Sequential(*layers)\nnet","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# Using Multiple Cloud TPU Cores (MultiProcessing)\n\nCommon issues:\nhttps://github.com/pytorch/xla/issues/1587\n\n## style 1"},{"metadata":{"trusted":true},"cell_type":"code","source":"def map_fn(index, flags):\n    \n    ## Sets a common random seed - both for initialization and ensuring graph is the same\n    torch.manual_seed(flags['seed'])\n    \n    ## Acquires the (unique) Cloud TPU core corresponding to this process's index\n    device = xm.xla_device()   \n    \n    ## Dataloader construction\n\n    train_dataset = mnistDataset(img_arrs,labels)\n    \n    # Creates the (distributed) sampler\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    # Creates dataloaders, which load data in batches\n    # Note: test loader is not shuffled or sampled (include in sampler)\n    train_loader = DataLoader(train_dataset,\n                              batch_size=flags['batch_size'],\n                              sampler=train_sampler,\n                              num_workers=flags['num_workers'],\n                              drop_last=True)\n    \n    \n    ## Network, optimizer, and loss function creation\n    \n    # create network Note: each process has its own identical copy of the model\n    net = torchvision.models.alexnet(num_classes=10)\n    layers = [nn.Conv2d(1, 64, kernel_size = 11, stride = 4, padding = 2)] + list(net.features.children())[1:]\n    net.features = nn.Sequential(*layers)\n    net = net.to(device).train()\n\n    loss_window = deque(maxlen=flags['window'])\n    \n    criterion = torch.nn.CrossEntropyLoss()\n    optimizer = torch.optim.Adam(net.parameters())\n    \n    ## Trains\n    tracker = xm.RateTracker()\n    train_start = time.time()\n    for epoch in range(flags['num_epochs']):\n        para_train_loader = pl.ParallelLoader(train_loader, [device]).per_device_loader(device)\n        for x, (data, targets) in enumerate(para_train_loader):\n            output = net(data)\n            loss = criterion(output, targets)\n            optimizer.zero_grad()\n            loss.backward()\n            loss_window.append(loss.item())\n\n      # Note: optimizer_step uses the implicit Cloud TPU context to\n      #  coordinate and synchronize gradient updates across processes.\n      #  This means that each process's network has the same weights after\n      #  this is called.\n        \n            xm.optimizer_step(optimizer)  # Note: barrier=True not needed when using ParallelLoader \n            tracker.add(flags['batch_size'])\n            if x % flags['window']  == 0:\n                print('[xla:{}]({}) Loss={:.5f} Rate={:.2f} GlobalRate={:.2f} Time={}'.format(\n            xm.get_ordinal(), x, np.mean(loss_window), tracker.rate(),\n            tracker.global_rate(), time.asctime()), flush=True)\n\n    elapsed_train_time = time.time() - train_start\n    print(\"Process\", index, \"finished training. Train time was:\", elapsed_train_time) \n","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"flags = {}\nflags['batch_size'] = 32\nflags['num_workers'] = 4\nflags['num_epochs'] = 1\nflags['seed'] = 1234\nflags['window'] = 20\n\nxmp.spawn(map_fn, args=(flags,), nprocs=8, start_method='fork')","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"## style 2"},{"metadata":{"trusted":true},"cell_type":"code","source":"def train_net():\n    torch.manual_seed(FLAGS['seed'])\n    device = xm.xla_device()   \n    train_dataset = mnistDataset(img_arrs,labels)\n    \n    train_sampler = DistributedSampler(train_dataset,\n                                       num_replicas = xm.xrt_world_size(),\n                                       rank = xm.get_ordinal(),\n                                       shuffle = True)\n    \n    train_loader = DataLoader(train_dataset,\n                              batch_size = FLAGS['batch_size'],\n                              sampler = train_sampler,\n                              num_workers = FLAGS['num_workers'],\n                              drop_last = True)\n    \n    ###\n    model = torchvision.models.alexnet(num_classes=10)\n    layers = [nn.Conv2d(1, 64, kernel_size = 11, stride = 4, padding = 2)] + list(net.features.children())[1:]\n    model.features = nn.Sequential(*layers)\n    model = net.to(device)\n    ###\n\n    loss_window = deque(maxlen = FLAGS['log_steps'])\n    loss_fn = torch.nn.CrossEntropyLoss()\n    optimizer = torch.optim.Adam(net.parameters())\n    \n    def train_loop_fn(loader):\n        tracker = xm.RateTracker()\n        model.train()\n        for x, (data, target) in enumerate(loader):\n            optimizer.zero_grad()\n            output = model(data)\n            loss = loss_fn(output, target)\n            loss.backward()\n            loss_window.append(loss.item())\n            xm.optimizer_step(optimizer)\n            tracker.add(FLAGS['batch_size'])\n            if (x+1) % FLAGS['log_steps'] == 0:\n                print('[xla:{}]({}) Loss={:.5f} Rate={:.2f} GlobalRate={:.2f} Time={}'.format(\n            xm.get_ordinal(), x+1, np.mean(loss_window), tracker.rate(),\n            tracker.global_rate(), time.asctime()), flush=True)\n    \n    for epoch in range(1, FLAGS['num_epochs'] + 1):\n        para_loader = pl.ParallelLoader(train_loader, [device])\n        train_loop_fn(para_loader.per_device_loader(device))\n        xm.master_print(\"Finished training epoch {}\".format(epoch))\n        \n    ","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"def _mp_fn(rank, flags):\n    global FLAGS\n    FLAGS = flags\n    torch.set_default_tensor_type('torch.FloatTensor')\n    train_start = time.time()\n    train_net()\n    elapsed_train_time = time.time() - train_start\n    print(\"Process\", rank, \"finished training. Train time was:\", elapsed_train_time)\n\n\n# Define Parameters\nFLAGS = {}\nFLAGS['seed'] = 1\nFLAGS['batch_size'] = 128\nFLAGS['num_workers'] = 4\nFLAGS['num_epochs'] = 2\nFLAGS['num_cores'] = 8\nFLAGS['log_steps'] = 20\nxmp.spawn(_mp_fn, args = (FLAGS,), nprocs = FLAGS['num_cores'],start_method='fork')","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# Using single CPU/GPU/TPU"},{"metadata":{"trusted":true},"cell_type":"code","source":"try:\n    device = xm.xla_device()\n    is_tpu = True\nexcept:\n    is_tpu = False\n    device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\nprint(device)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"net = torchvision.models.alexnet(num_classes=10)\nlayers = [nn.Conv2d(1, 64, kernel_size = 11, stride = 4, padding = 2)] + list(net.features.children())[1:]\nnet.features = nn.Sequential(*layers)\nnet = net.to(device).train()\n\ncriterion = torch.nn.CrossEntropyLoss()\noptimizer = torch.optim.Adam(net.parameters())\nnum_epochs = 1\nwindow = 100\nloss_window = deque(maxlen=window)\n\nstart_time = time.time()\n\nfor epoch in range(num_epochs):\n    for counter, (data, targets) in enumerate(train_loader):\n        data = data.to(device)\n        targets = targets.to(device)\n        output = net(data)\n\n        loss = criterion(output, targets)\n        loss_window.append(loss.item())\n\n        optimizer.zero_grad()\n        loss.backward()\n        if is_tpu: \n            xm.optimizer_step(optimizer, barrier=True)  # Note: Cloud TPU-specific code!\n        else:\n            optimizer.step()\n            \n        print('\\rEpoch {}/{}\\tAverage Score: {:.2f}'.format(epoch,counter*batch_size, np.mean(loss_window)),end=\"\")\n        if counter % window == 0:\n            print('\\rEpoch {}/{}\\tAverage Score: {:.2f}'.format(epoch,counter*batch_size, np.mean(loss_window)))\n\nelapsed_time = time.time() - start_time\nprint (\"\\t Spent \", elapsed_time, \" seconds training for \", num_epochs, \" epoch(s) on a single core.\")","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}