codes/MMEval_tutorials/2-MMEval-使用.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "f03f4492",
   "metadata": {},
   "source": [
    "本节将以 CIFAR-10 数据集的评测为例，分别介绍如何使用 MMEval 结合 torch.distributed 和 MPI4Py 进行分布式评测，相关代码可以在 [mmeval/examples/cifar10_dist_eval](https://github.com/open-mmlab/mmeval/tree/main/examples/cifar10_dist_eval) 中找到。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "9b4c32f3",
   "metadata": {
    "scrolled": false
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Looking in indexes: https://pypi.tuna.tsinghua.edu.cn/simple\n",
      "Requirement already satisfied: torch in /nvme/data/xiaoyancong/miniconda3/envs/mmeval/lib/python3.8/site-packages (1.12.1)\n",
      "Requirement already satisfied: torchvision in /nvme/data/xiaoyancong/miniconda3/envs/mmeval/lib/python3.8/site-packages (0.13.1)\n",
      "Requirement already satisfied: tqdm in /nvme/data/xiaoyancong/miniconda3/envs/mmeval/lib/python3.8/site-packages (4.64.1)\n",
      "Requirement already satisfied: typing_extensions in /nvme/data/xiaoyancong/miniconda3/envs/mmeval/lib/python3.8/site-packages (from torch) (4.4.0)\n",
      "Requirement already satisfied: numpy in /nvme/data/xiaoyancong/miniconda3/envs/mmeval/lib/python3.8/site-packages (from torchvision) (1.23.3)\n",
      "Requirement already satisfied: requests in /nvme/data/xiaoyancong/miniconda3/envs/mmeval/lib/python3.8/site-packages (from torchvision) (2.28.1)\n",
      "Requirement already satisfied: pillow!=8.3.*,>=5.3.0 in /nvme/data/xiaoyancong/miniconda3/envs/mmeval/lib/python3.8/site-packages (from torchvision) (9.2.0)\n",
      "Requirement already satisfied: certifi>=2017.4.17 in /nvme/data/xiaoyancong/miniconda3/envs/mmeval/lib/python3.8/site-packages (from requests->torchvision) (2022.9.24)\n",
      "Requirement already satisfied: charset-normalizer<3,>=2 in /nvme/data/xiaoyancong/miniconda3/envs/mmeval/lib/python3.8/site-packages (from requests->torchvision) (2.0.4)\n",
      "Requirement already satisfied: urllib3<1.27,>=1.21.1 in /nvme/data/xiaoyancong/miniconda3/envs/mmeval/lib/python3.8/site-packages (from requests->torchvision) (1.26.12)\n",
      "Requirement already satisfied: idna<4,>=2.5 in /nvme/data/xiaoyancong/miniconda3/envs/mmeval/lib/python3.8/site-packages (from requests->torchvision) (3.4)\n"
     ]
    }
   ],
   "source": [
    "!pip install torch torchvision tqdm"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "9ed6cb08",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "2022-11-16 22:05:56.763074: I tensorflow/stream_executor/platform/default/dso_loader.cc:49] Successfully opened dynamic library libcudart.so.11.0\n"
     ]
    }
   ],
   "source": [
    "import torch\n",
    "import torchvision as tv\n",
    "import tqdm\n",
    "from torch.utils.data import DataLoader\n",
    "\n",
    "from mmeval import Accuracy"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7f696518",
   "metadata": {},
   "source": [
    "## 1. 单进程评测"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e4598813",
   "metadata": {},
   "source": [
    "首先我们需要加载 CIFAR-10 测试数据，我们可以使用 TorchVison 提供的数据集类。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "dfcb3567",
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_eval_dataloader():\n",
    "    dataset = tv.datasets.CIFAR10(\n",
    "        root='./',\n",
    "        train=False,\n",
    "        download=True,\n",
    "        transform=tv.transforms.ToTensor())\n",
    "    return DataLoader(dataset, batch_size=1)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5b548525",
   "metadata": {},
   "source": [
    "其次，我们需要准备待评测的模型，这里我们使用 TorchVision 中的 resnet18。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "54fc3c57",
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_model(pretrained_model_fpath=None):\n",
    "    model = tv.models.resnet18(num_classes=10)\n",
    "    if pretrained_model_fpath is not None:\n",
    "        model.load_state_dict(torch.load(pretrained_model_fpath))\n",
    "    return model.eval()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f1a22235",
   "metadata": {},
   "source": [
    "有了待评测的数据集与模型，就可以使用 mmeval.Accuracy 指标对模型预测结果进行评测。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "d6b2097e",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Files already downloaded and verified\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|███████████████████████████████████████████████████████████████████████████| 10000/10000 [00:33<00:00, 302.07it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'top1': 0.7458999752998352, 'top3': 0.8931000232696533}\n"
     ]
    }
   ],
   "source": [
    "eval_dataloader = get_eval_dataloader()\n",
    "model = get_model('./cifar10_resnet18.pth').cuda()\n",
    "# 实例化 `Accuracy`，计算 top1 与 top3 准确率\n",
    "accuracy = Accuracy(topk=(1, 3))\n",
    "\n",
    "with torch.no_grad():\n",
    "    for images, labels in tqdm.tqdm(eval_dataloader):\n",
    "        predicted_score = model(images.cuda()).cpu()\n",
    "        # 累计批次数据，中间结果将保存在 `accuracy._results` 中\n",
    "        accuracy.add(predictions=predicted_score, labels=labels)\n",
    "\n",
    "# 调用 `accuracy.compute` 进行指标计算\n",
    "print(accuracy.compute())\n",
    "# 调用 `accuracy.reset` 清除保存在 `accuracy._results` 中的中间结果\n",
    "accuracy.reset()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5f98e385",
   "metadata": {},
   "source": [
    "## 2. 使用 torch.distributed 进行分布式评测"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2e819510",
   "metadata": {},
   "source": [
    "在 MMEval 中为 torch.distributed 实现了两个分布式通信后端，分别是 TorchCPUDist 和 TorchCUDADist。\n",
    "\n",
    "为 MMEval 设置分布式通信后端的方式有两种："
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "ec104453",
   "metadata": {},
   "outputs": [],
   "source": [
    "from mmeval.core import set_default_dist_backend\n",
    "from mmeval import Accuracy\n",
    "\n",
    "# 1. 设置全局默认分布式通信后端\n",
    "set_default_dist_backend('torch_cpu')\n",
    "\n",
    "# 2. 初始化评测指标时候通过 `dist_backend` 传参\n",
    "accuracy = Accuracy(dist_backend='torch_cpu')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a499999a",
   "metadata": {},
   "source": [
    "结合上述单进程评测的代码，再加入数据集切片和分布式初始化，即可实现分布式评测。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "94e921fc",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "import torch\r\n",
      "import torchvision as tv\r\n",
      "import tqdm\r\n",
      "from torch.utils.data import DataLoader, DistributedSampler\r\n",
      "\r\n",
      "from mmeval import Accuracy\r\n",
      "\r\n",
      "\r\n",
      "def get_eval_dataloader(rank=0, num_replicas=1):\r\n",
      "    dataset = tv.datasets.CIFAR10(\r\n",
      "        root='./',\r\n",
      "        train=False,\r\n",
      "        download=True,\r\n",
      "        transform=tv.transforms.ToTensor())\r\n",
      "    dist_sampler = DistributedSampler(\r\n",
      "        dataset, num_replicas=num_replicas, rank=rank)\r\n",
      "    data_loader = DataLoader(dataset, batch_size=1, sampler=dist_sampler)\r\n",
      "    return data_loader, len(dataset)\r\n",
      "\r\n",
      "\r\n",
      "def get_model(pretrained_model_fpath=None):\r\n",
      "    model = tv.models.resnet18(num_classes=10)\r\n",
      "    if pretrained_model_fpath is not None:\r\n",
      "        model.load_state_dict(torch.load(pretrained_model_fpath))\r\n",
      "    return model.eval()\r\n",
      "\r\n",
      "\r\n",
      "def eval_fn(rank, process_num):\r\n",
      "    torch.distributed.init_process_group(\r\n",
      "        backend='gloo',\r\n",
      "        init_method='tcp://127.0.0.1:2345',\r\n",
      "        world_size=process_num,\r\n",
      "        rank=rank)\r\n",
      "    torch.cuda.set_device(rank)\r\n",
      "\r\n",
      "    eval_dataloader, total_num_samples = get_eval_dataloader(rank, process_num)\r\n",
      "    model = get_model('./cifar10_resnet18.pth').cuda()\r\n",
      "    accuracy = Accuracy(topk=(1, 3), dist_backend='torch_cpu')\r\n",
      "\r\n",
      "    with torch.no_grad():\r\n",
      "        for images, labels in tqdm.tqdm(eval_dataloader, disable=(rank != 0)):\r\n",
      "            predicted_score = model(images.cuda()).cpu()\r\n",
      "            accuracy.add(predictions=predicted_score, labels=labels)\r\n",
      "\r\n",
      "    print(accuracy.compute(size=total_num_samples))\r\n",
      "    accuracy.reset()\r\n",
      "\r\n",
      "\r\n",
      "if __name__ == '__main__':\r\n",
      "    process_num = 3\r\n",
      "    torch.multiprocessing.spawn(\r\n",
      "        eval_fn, nprocs=process_num, args=(process_num, ))\r\n"
     ]
    }
   ],
   "source": [
    "!cat cifar10_dist_eval/cifar10_eval_torch_dist.py"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "2701a60b",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "2022-11-16 22:10:49.905219: I tensorflow/stream_executor/platform/default/dso_loader.cc:49] Successfully opened dynamic library libcudart.so.11.0\n",
      "2022-11-16 22:10:53.194848: I tensorflow/stream_executor/platform/default/dso_loader.cc:49] Successfully opened dynamic library libcudart.so.11.0\n",
      "2022-11-16 22:10:53.194848: I tensorflow/stream_executor/platform/default/dso_loader.cc:49] Successfully opened dynamic library libcudart.so.11.0\n",
      "2022-11-16 22:10:53.207048: I tensorflow/stream_executor/platform/default/dso_loader.cc:49] Successfully opened dynamic library libcudart.so.11.0\n",
      "Files already downloaded and verified\n",
      "Files already downloaded and verified\n",
      "Files already downloaded and verified\n",
      "100%|██████████████████████████████████████| 3334/3334 [00:11<00:00, 283.87it/s]\n",
      "{'top1': 0.7458999752998352, 'top3': 0.8931000232696533}\n",
      "{'top1': 0.7458999752998352, 'top3': 0.8931000232696533}\n",
      "{'top1': 0.7458999752998352, 'top3': 0.8931000232696533}\n"
     ]
    }
   ],
   "source": [
    "!python cifar10_dist_eval/cifar10_eval_torch_dist.py"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "68763bff",
   "metadata": {},
   "source": [
    "## 3. 使用 MPI4Py 进行分布式评测"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6b26bc33",
   "metadata": {},
   "source": [
    "MMEval 将分布式通信功能抽象解耦了，因此虽然上述例子使用的是 PyTorch 模型和数据加载，我们仍然可以使用除 torch.distributed 以外的分布式通信后端来实现分布式评测。下面将展示如何使用 MPI4Py 作为分布式通信后端来进行分布式评测。"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4ede318a",
   "metadata": {},
   "source": [
    "首先需要安装 MPI4Py 以及 openmpi，建议使用 conda 进行安装："
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "9398f9bf",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Collecting package metadata (current_repodata.json): done\n",
      "Solving environment: done\n",
      "\n",
      "# All requested packages already installed.\n",
      "\n"
     ]
    }
   ],
   "source": [
    "!conda install -y openmpi mpi4py"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "9a649abf",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "import torch\r\n",
      "import torchvision as tv\r\n",
      "import tqdm\r\n",
      "from mpi4py import MPI\r\n",
      "from torch.utils.data import DataLoader, DistributedSampler\r\n",
      "\r\n",
      "from mmeval import Accuracy\r\n",
      "\r\n",
      "\r\n",
      "def get_eval_dataloader(rank=0, num_replicas=1):\r\n",
      "    dataset = tv.datasets.CIFAR10(\r\n",
      "        root='./',\r\n",
      "        train=False,\r\n",
      "        download=True,\r\n",
      "        transform=tv.transforms.ToTensor())\r\n",
      "    dist_sampler = DistributedSampler(\r\n",
      "        dataset, num_replicas=num_replicas, rank=rank)\r\n",
      "    data_loader = DataLoader(dataset, batch_size=1, sampler=dist_sampler)\r\n",
      "    return data_loader, len(dataset)\r\n",
      "\r\n",
      "\r\n",
      "def get_model(pretrained_model_fpath=None):\r\n",
      "    model = tv.models.resnet18(num_classes=10)\r\n",
      "    if pretrained_model_fpath is not None:\r\n",
      "        model.load_state_dict(torch.load(pretrained_model_fpath))\r\n",
      "    return model.eval()\r\n",
      "\r\n",
      "\r\n",
      "def eval_fn(rank, process_num):\r\n",
      "    torch.cuda.set_device(rank)\r\n",
      "    eval_dataloader, total_num_samples = get_eval_dataloader(rank, process_num)\r\n",
      "    model = get_model('./cifar10_resnet18.pth').cuda()\r\n",
      "    accuracy = Accuracy(topk=(1, 3), dist_backend='mpi4py')\r\n",
      "\r\n",
      "    with torch.no_grad():\r\n",
      "        for images, labels in tqdm.tqdm(eval_dataloader, disable=(rank != 0)):\r\n",
      "            predicted_score = model(images.cuda()).cpu()\r\n",
      "            accuracy.add(predictions=predicted_score, labels=labels)\r\n",
      "\r\n",
      "    print(accuracy.compute(size=total_num_samples))\r\n",
      "    accuracy.reset()\r\n",
      "\r\n",
      "\r\n",
      "if __name__ == '__main__':\r\n",
      "    comm = MPI.COMM_WORLD\r\n",
      "    eval_fn(comm.Get_rank(), comm.Get_size())\r\n"
     ]
    }
   ],
   "source": [
    "!cat cifar10_dist_eval/cifar10_eval_mpi4py.py"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6f585b4c",
   "metadata": {},
   "source": [
    "使用 mpirun 作为分布式评测启动方式："
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "843293fa",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "2022-11-16 22:12:59.873751: I tensorflow/stream_executor/platform/default/dso_loader.cc:49] Successfully opened dynamic library libcudart.so.11.0\n",
      "2022-11-16 22:12:59.873752: I tensorflow/stream_executor/platform/default/dso_loader.cc:49] Successfully opened dynamic library libcudart.so.11.0\n",
      "2022-11-16 22:12:59.874402: I tensorflow/stream_executor/platform/default/dso_loader.cc:49] Successfully opened dynamic library libcudart.so.11.0\n",
      "Files already downloaded and verified\n",
      "Files already downloaded and verified\n",
      "Files already downloaded and verified\n",
      "100%|██████████| 3334/3334 [00:11<00:00, 282.08it/s]{'top1': 0.7458999752998352, 'top3': 0.8931000232696533}\n",
      "{'top1': 0.7458999752998352, 'top3': 0.8931000232696533}\n",
      "{'top1': 0.7458999752998352, 'top3': 0.8931000232696533}\n",
      "\n"
     ]
    }
   ],
   "source": [
    "# 使用 mpirun 启动 3 个进程\n",
    "!mpirun -np 3 python cifar10_dist_eval/cifar10_eval_mpi4py.py"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d5ea95e7",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "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.8.13"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}