Added XPU example for non-federated_case

This commit introduces an example of using an XPU with the non-federated_case.

openfl-tutorials/interactive_api/PyTorch_TinyImageNet/workspace/non-federated_case_XPU.ipynb

@@ -0,0 +1,357 @@
+{
+ "cells": [
+  {
+   "cell_type": "markdown",
+   "id": "4d452110-525d-47d8-b9fc-f3173cfd2733",
+   "metadata": {},
+   "source": [
+    "## Vanilla PyTorch training on TinyImageNet dataset (XPU version)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "a20dc008-7c65-4ab4-9fb4-39605bb5599d",
+   "metadata": {},
+   "source": [
+    "This notebook is intended to show that fixing random seeds leads to the same result in both federated and non-federated cases. \n",
+    "\n",
+    "Also, as a simple example for use Intel® Extension for PyTorch* with the latest performance optimizations for Intel hardware. Intel® Extension for PyTorch* provides easy GPU acceleration for Intel discrete GPUs through the PyTorch* xpu device.\r\n",
+    "\n",
+    "*Please refer to the [Installation Guide](https://intel.github.io/intel-extension-for-pytorch/xpu/2.1.10+xpu/tutorials/installation.html) for the system requirements and steps to install and use Intel® Extension for PyTorch*. For more detailed tutorials and documentations describing features, APIs and technical details, please refer to [Intel® Extension for PyTorch* Documentation](https://intel.github.io/intel-extension-for-pytorch/xpu/2.1.10+xpu/index.html). "
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "e2a00108-cdb0-47ff-8ed5-71d92bcc20be",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# Check Installation guide and Documentacion of IPEX for additional requiriments\n",
+    "!pip install -r requirements.txt"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "bb1add50-c2c9-4881-a4a7-9bff634d219f",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "from pathlib import Path\n",
+    "import os\n",
+    "import shutil\n",
+    "from torch.utils.data import Dataset\n",
+    "from torch.utils.data import DataLoader\n",
+    "from torch import nn\n",
+    "from torch import optim\n",
+    "import torch.nn.functional as F\n",
+    "import torch\n",
+    "import torchvision.transforms as T\n",
+    "import torchvision\n",
+    "import glob\n",
+    "import tqdm\n",
+    "from PIL import Image\n",
+    "import numpy as np\n",
+    "\n",
+    "############# code changes ###############\n",
+    "device='xpu'"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "9a365362-2224-41c8-b01d-c76229331d6f",
+   "metadata": {},
+   "source": [
+    "## Download data"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "fcc1cb4e-baca-4380-9cb8-f9a96e985779",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "common_data_folder = Path.cwd() / 'data'\n",
+    "zip_file_path = common_data_folder / 'tiny-imagenet-200.zip'\n",
+    "os.makedirs(common_data_folder, exist_ok=True)\n",
+    "os.system(f'wget --no-clobber http://cs231n.stanford.edu/tiny-imagenet-200.zip'\n",
+    "          f' -O {zip_file_path}')\n",
+    "shutil.unpack_archive(str(zip_file_path), str(common_data_folder))"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "fdd48081-0db6-467a-9b1d-b66aaf8bedbb",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "class TinyImageNetDataset(Dataset):\n",
+    "    \"\"\"TinyImageNet shard dataset class.\"\"\"\n",
+    "\n",
+    "    NUM_IMAGES_PER_CLASS = 500\n",
+    "\n",
+    "    def __init__(self, data_folder: Path, data_type='train', transform=None):\n",
+    "        \"\"\"Initialize TinyImageNetDataset.\"\"\"\n",
+    "        self.data_type = data_type\n",
+    "        self._common_data_folder = data_folder\n",
+    "        self._data_folder = os.path.join(data_folder, data_type)\n",
+    "        self.labels = {}  # fname - label number mapping\n",
+    "        self.image_paths = sorted(\n",
+    "            glob.iglob(\n",
+    "                os.path.join(self._data_folder, '**', '*.JPEG'),\n",
+    "                recursive=True\n",
+    "            )\n",
+    "        )\n",
+    "        with open(os.path.join(self._common_data_folder, 'wnids.txt'), 'r') as fp:\n",
+    "            self.label_texts = sorted([text.strip() for text in fp.readlines()])\n",
+    "        self.label_text_to_number = {text: i for i, text in enumerate(self.label_texts)}\n",
+    "        self.fill_labels()\n",
+    "        self.transform = transform\n",
+    "\n",
+    "    def __len__(self) -> int:\n",
+    "        \"\"\"Return the len of the shard dataset.\"\"\"\n",
+    "        return len(self.image_paths)\n",
+    "\n",
+    "    def __getitem__(self, index: int):\n",
+    "        \"\"\"Return an item by the index.\"\"\"\n",
+    "        file_path = self.image_paths[index]\n",
+    "        sample = self.read_image(file_path)\n",
+    "        if self.transform:\n",
+    "            sample = self.transform(sample)\n",
+    "        label = self.labels[os.path.basename(file_path)]\n",
+    "        return sample, label\n",
+    "\n",
+    "    def read_image(self, path: Path):\n",
+    "        \"\"\"Read the image.\"\"\"\n",
+    "        img = Image.open(path)\n",
+    "        return img\n",
+    "\n",
+    "    def fill_labels(self) -> None:\n",
+    "        \"\"\"Fill labels.\"\"\"\n",
+    "        if self.data_type == 'train':\n",
+    "            for label_text, i in self.label_text_to_number.items():\n",
+    "                for cnt in range(self.NUM_IMAGES_PER_CLASS):\n",
+    "                    self.labels[f'{label_text}_{cnt}.JPEG'] = i\n",
+    "        elif self.data_type == 'val':\n",
+    "            with open(os.path.join(self._data_folder, 'val_annotations.txt'), 'r') as fp:\n",
+    "                for line in fp.readlines():\n",
+    "                    terms = line.split('\\t')\n",
+    "                    file_name, label_text = terms[0], terms[1]\n",
+    "                    self.labels[file_name] = self.label_text_to_number[label_text]"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "5f08dbff-08e8-4ddc-86d9-6d20e510affd",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "normalize = T.Normalize(\n",
+    "    mean=[0.485, 0.456, 0.406],\n",
+    "    std=[0.229, 0.224, 0.225]\n",
+    ")\n",
+    "\n",
+    "augmentation = T.RandomApply(\n",
+    "    [T.RandomHorizontalFlip(),\n",
+    "     T.RandomRotation(10),\n",
+    "     T.RandomResizedCrop(64)], \n",
+    "    p=.8\n",
+    ")\n",
+    "\n",
+    "training_transform = T.Compose(\n",
+    "    [T.Lambda(lambda x: x.convert(\"RGB\")),\n",
+    "     T.ToTensor(),\n",
+    "     augmentation,\n",
+    "     normalize]\n",
+    ")\n",
+    "\n",
+    "valid_transform = T.Compose(\n",
+    "    [T.Lambda(lambda x: x.convert(\"RGB\")),\n",
+    "     T.ToTensor(),\n",
+    "     normalize]\n",
+    ")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "edb570c5-cd8c-4aa3-be52-43b579ea13b8",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "def get_train_loader():\n",
+    "    generator=torch.Generator()\n",
+    "    generator.manual_seed(0)\n",
+    "    train_set = TinyImageNetDataset(common_data_folder / 'tiny-imagenet-200', transform=training_transform)\n",
+    "    return DataLoader(train_set, batch_size=64, shuffle=True, generator=generator)\n",
+    "\n",
+    "def get_valid_loader():\n",
+    "    valid_set = TinyImageNetDataset(common_data_folder / 'tiny-imagenet-200', data_type='val', transform=valid_transform)\n",
+    "    return DataLoader(valid_set, batch_size=64)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "4fd157ac-0f37-42a7-9f79-48e6b4e6b27e",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "class Net(nn.Module):\n",
+    "    def __init__(self):\n",
+    "        torch.manual_seed(0)\n",
+    "        super(Net, self).__init__()\n",
+    "        self.model = torchvision.models.mobilenet_v2(pretrained=True)\n",
+    "        # self.model.requires_grad_(False)\n",
+    "        self.model.classifier[1] = torch.nn.Linear(in_features=1280, \\\n",
+    "                        out_features=200, bias=True)\n",
+    "\n",
+    "    def forward(self, x):\n",
+    "        x = self.model.forward(x)\n",
+    "        return x\n",
+    "\n",
+    "model = Net()"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "99377a32-9867-49d0-b17c-83f22ad88ecc",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "optimizer = optim.Adam([x for x in model.parameters() if x.requires_grad], lr=1e-4)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "c4aac8d0-1a12-40b1-a1e2-59cc82cea3f9",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "loss_fn = F.cross_entropy"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "4d4bb282-7231-47a1-9342-0d44e4f76d98",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "############# code changes ###############\n",
+    "import intel_extension_for_pytorch as ipex\n",
+    "model.to(device)\n",
+    "model, optimizer = ipex.optimize(model, optimizer=optimizer, dtype=torch.bfloat16)\n",
+    "\n",
+    "def train(model, optimizer):\n",
+    "    torch.manual_seed(0)\n",
+    "    \n",
+    "    data_loader = tqdm.tqdm(get_train_loader(), desc=\"train\")\n",
+    "    model.train()\n",
+    "    losses = []\n",
+    "\n",
+    "    for data, target in data_loader:\n",
+    "        ############# code changes ###############\n",
+    "        data, target = data.to(device), target.to(device, dtype=torch.int64)\n",
+    "        #data, target = torch.tensor(data).to(device), torch.tensor(\n",
+    "        #    target).to(device)\n",
+    "        optimizer.zero_grad()\n",
+    "        ############# code changes ###############\n",
+    "        with torch.xpu.amp.autocast(enabled=True, dtype=torch.bfloat16):\n",
+    "            output = model(data)\n",
+    "            loss = loss_fn(output, target)\n",
+    "        loss.backward()\n",
+    "        optimizer.step()\n",
+    "        losses.append(loss.detach().cpu().numpy())\n",
+    "        \n",
+    "    return {'train_loss': np.mean(losses),}\n",
+    "\n",
+    "def validate(model):\n",
+    "    torch.manual_seed(0)\n",
+    "    model.eval()\n",
+    "    \n",
+    "    data_loader = tqdm.tqdm(get_valid_loader(), desc=\"validate\")\n",
+    "    val_score = 0\n",
+    "    total_samples = 0\n",
+    "\n",
+    "    with torch.no_grad():\n",
+    "        for data, target in data_loader:\n",
+    "            samples = target.shape[0]\n",
+    "            total_samples += samples\n",
+    "            ############# code changes ###############\n",
+    "            data, target = data.to(device), target.to(device, dtype=torch.int64)\n",
+    "            #data, target = torch.tensor(data).to(device), \\\n",
+    "            #    torch.tensor(target).to(device, dtype=torch.int64)\n",
+    "            ############# code changes ###############\n",
+    "            with torch.xpu.amp.autocast(enabled=True, dtype=torch.bfloat16):\n",
+    "                output = model(data)\n",
+    "                pred = output.argmax(dim=1,keepdim=True)\n",
+    "            val_score += pred.eq(target).sum().cpu().numpy()\n",
+    "            \n",
+    "    return {'acc': val_score / total_samples,}"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "0baeafef-212b-4ca8-8edc-baa4904d4b08",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "for i in range(5):\n",
+    "    if i == 0:\n",
+    "        name, value = next(iter(validate(model).items())) \n",
+    "        print(f'{name}: {value:f}')\n",
+    "    \n",
+    "    name, value = next(iter(train(model, optimizer).items()))\n",
+    "    print(f'{name}: {value:f}')\n",
+    "    \n",
+    "    name, value = next(iter(validate(model).items())) \n",
+    "    print(f'{name}: {value:f}')"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "ed198e40-453a-4a52-8905-665744750c7e",
+   "metadata": {},
+   "outputs": [],
+   "source": []
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "b5fd54de-93d9-4440-b7f1-d20e22a4fc71",
+   "metadata": {},
+   "outputs": [],
+   "source": []
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "open_fl",
+   "language": "python",
+   "name": "open_fl"
+  },
+  "language_info": {
+   "codemirror_mode": {
+    "name": "ipython",
+    "version": 3
+   },
+   "file_extension": ".py",
+   "mimetype": "text/x-python",
+   "name": "python",
+   "nbconvert_exporter": "python",
+   "pygments_lexer": "ipython3",
+   "version": "3.9.18"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 5
+}