create_model.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\u001b[38;5;2m[i 1227 09:44:38.842422 48 compiler.py:944] Jittor(1.3.1.27) src: /home/lry/miniconda3/lib/python3.8/site-packages/jittor\u001b[m\n",
      "\u001b[38;5;2m[i 1227 09:44:38.853090 48 compiler.py:945] g++ at /usr/bin/g++(7.5.0)\u001b[m\n",
      "\u001b[38;5;2m[i 1227 09:44:38.854896 48 compiler.py:946] cache_path: /home/lry/.cache/jittor/jt1.3.1/g++7.5.0/py3.8.5/Linux-4.15.0-1xe3/IntelRXeonRSilxdc/default\u001b[m\n",
      "\u001b[38;5;2m[i 1227 09:44:38.865654 48 __init__.py:372] Found nvcc(11.1.105) at /usr/local/cuda/bin/nvcc.\u001b[m\n",
      "\u001b[38;5;2m[i 1227 09:44:38.877402 48 __init__.py:372] Found addr2line(2.30) at /usr/bin/addr2line.\u001b[m\n",
      "\u001b[38;5;2m[i 1227 09:44:39.073720 48 compiler.py:997] cuda key:cu11.1.105_sm_86\u001b[m\n",
      "\u001b[38;5;2m[i 1227 09:44:39.321505 48 __init__.py:187] Total mem: 125.56GB, using 16 procs for compiling.\u001b[m\n",
      "\u001b[38;5;2m[i 1227 09:44:39.474608 48 jit_compiler.cc:27] Load cc_path: /usr/bin/g++\u001b[m\n",
      "\u001b[38;5;2m[i 1227 09:44:39.623317 48 init.cc:61] Found cuda archs: [86,]\u001b[m\n",
      "\u001b[38;5;2m[i 1227 09:44:39.815691 48 __init__.py:372] Found mpicc(2.1.1) at /usr/bin/mpicc.\u001b[m\n",
      "\u001b[38;5;2m[i 1227 09:44:40.018778 48 compile_extern.py:29] found /usr/local/cuda/include/cublas.h\u001b[m\n",
      "\u001b[38;5;2m[i 1227 09:44:40.037088 48 compile_extern.py:29] found /usr/local/cuda/lib64/libcublas.so\u001b[m\n",
      "\u001b[38;5;2m[i 1227 09:44:40.039368 48 compile_extern.py:29] found /usr/local/cuda/lib64/libcublasLt.so.11\u001b[m\n",
      "\u001b[38;5;2m[i 1227 09:44:40.871322 48 compile_extern.py:29] found /usr/local/cuda/include/cudnn.h\u001b[m\n",
      "\u001b[38;5;2m[i 1227 09:44:40.905350 48 compile_extern.py:29] found /usr/local/cuda/lib64/libcudnn.so.8\u001b[m\n",
      "\u001b[38;5;2m[i 1227 09:44:40.907306 48 compile_extern.py:29] found /usr/local/cuda/lib64/libcudnn_ops_infer.so.8\u001b[m\n",
      "\u001b[38;5;2m[i 1227 09:44:40.913792 48 compile_extern.py:29] found /usr/local/cuda/lib64/libcudnn_ops_train.so.8\u001b[m\n",
      "\u001b[38;5;2m[i 1227 09:44:40.915126 48 compile_extern.py:29] found /usr/local/cuda/lib64/libcudnn_cnn_infer.so.8\u001b[m\n",
      "\u001b[38;5;2m[i 1227 09:44:41.148306 48 compile_extern.py:29] found /usr/local/cuda/lib64/libcudnn_cnn_train.so.8\u001b[m\n",
      "\u001b[38;5;2m[i 1227 09:44:41.900577 48 compile_extern.py:29] found /usr/local/cuda/include/curand.h\u001b[m\n",
      "\u001b[38;5;2m[i 1227 09:44:41.952581 48 compile_extern.py:29] found /usr/local/cuda/lib64/libcurand.so\u001b[m\n"
     ]
    }
   ],
   "source": [
    "import jittor as jt\n",
    "from models_jittor import gMLPForImageClassification as gMLP_jt\n",
    "from models_jittor import ResMLPForImageClassification as ResMLP_jt\n",
    "from models_jittor import MLPMixerForImageClassification as MLPMixer_jt\n",
    "from models_jittor import ViP as ViP_jt\n",
    "from models_jittor import S2MLPv2 as S2MLPv2_jt\n",
    "from models_jittor import ConvMixer as ConvMixer_jt\n",
    "from models_jittor import convmlp_s as ConvMLP_s_jt \n",
    "from models_jittor import convmlp_l as ConvMLP_l_jt \n",
    "from models_jittor import convmlp_m as ConvMLP_m_jt \n",
    "from models_jittor import RaftMLP as RaftMLP_jt\n",
    "from models_jittor import SparseMLP as SparseMLP_jt\n",
    "from models_jittor import HireMLP as HireMLP_jt\n",
    "from models_jittor import AS_MLP as AS_MLP_jt\n",
    "from models_jittor import S2MLPv1_deep, S2MLPv1_wide\n",
    "from models_jittor import SwinMLP as SwinMLP_jt\n",
    "from models_jittor import WaveMLP as WaveMLP_jt\n",
    "from models_jittor import MS_MLP as MS_MLP_jt\n",
    "from models_jittor import MorphMLP as MorphMLP_jt\n",
    "from models_jittor import DynaMixer as DynaMixer_jt\n",
    "from models_jittor import Sequencer2D as Sequencer2D_jt"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "from models_pytorch import gMLPForImageClassification as gMLP_pt\n",
    "from models_pytorch import ResMLPForImageClassification as ResMLP_pt\n",
    "from models_pytorch import MLPMixerForImageClassification as MLPMixer_pt\n",
    "from models_pytorch import ViP as ViP_pt\n",
    "from models_pytorch import S2MLPv2 as S2MLPv2_pt \n",
    "from models_pytorch import ConvMixer as ConvMixer_pt \n",
    "from models_pytorch import convmlp_s as ConvMLP_s_pt \n",
    "from models_pytorch import convmlp_l as ConvMLP_l_pt \n",
    "from models_pytorch import convmlp_m as ConvMLP_m_pt \n",
    "from models_pytorch import RaftMLP as RaftMLP_pt\n",
    "from models_pytorch import SparseMLP as SparseMLP_pt\n",
    "from models_pytorch import HireMLP as HireMLP_pt\n",
    "from models_pytorch import GFNet as GFNet_pt\n",
    "from models_pytorch import CycleMLP_B2 as CycleMLP_B2_pt\n",
    "from models_pytorch import AS_MLP as AS_MLP_pt\n",
    "from models_pytorch import SwinMLP as SwinMLP_pt\n",
    "from models_pytorch import create_RepMLPNet_B224, create_RepMLPNet_B256\n",
    "from models_pytorch import WaveMLP as WaveMLP_pt\n",
    "from models_pytorch import MS_MLP as MS_MLP_pt\n",
    "from models_pytorch import MorphMLP as MorphMLP_pt\n",
    "from models_pytorch import DynaMixer as DynaMixer_pt\n",
    "from models_pytorch import Sequencer2D as Sequencer2D_pt\n",
    "import time"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "# model_jt = gMLP_jt(\n",
    "#     image_size=224,\n",
    "#     patch_size=16,\n",
    "#     in_channels=3,\n",
    "#     num_classes=1000,\n",
    "#     d_model=256,\n",
    "#     d_ffn=1536,\n",
    "#     depth=30,\n",
    "# )\n",
    "\n",
    "# model_jt = ResMLP_jt(\n",
    "#     in_channels = 3,\n",
    "#     image_size = 224,\n",
    "#     patch_size = 16,\n",
    "#     d_model = 384,\n",
    "#     depth = 12,\n",
    "#     num_classes = 1000, \n",
    "#     expansion_factor = 4\n",
    "# )\n",
    "\n",
    "# model_jt = MLPMixer_jt(\n",
    "#     image_size=(224,112),\n",
    "#     patch_size=(16, 8),\n",
    "#     in_channels=3,\n",
    "#     num_classes=1000,\n",
    "#     d_model=256,\n",
    "#     depth=12,\n",
    "# )\n",
    "\n",
    "# model_jt = ViP_jt(\n",
    "#     image_size=(224, 224),\n",
    "#     patch_size=(16, 8),\n",
    "#     in_channels=3,\n",
    "#     num_classes=1000,\n",
    "#     d_model=256,\n",
    "#     depth=30,\n",
    "#     segments = 16,\n",
    "#     expansion_factor = 4,\n",
    "#     weighted = False\n",
    "# )\n",
    "\n",
    "# model_jt = S2MLPv2_jt(\n",
    "#     in_channels = 3,\n",
    "#     image_size = (224,224),\n",
    "#     patch_size = [7],\n",
    "#     d_model = [384],\n",
    "#     depth = [18],\n",
    "#     num_classes = 1000, \n",
    "#     expansion_factor = [3]\n",
    "# )\n",
    "\n",
    "# model_jt = ConvMixer_jt(\n",
    "#     dim = 1568,\n",
    "#     depth = 20\n",
    "# )\n",
    "\n",
    "# model_jt = ConvMLP_s_jt(pretrained = True, num_classes = 1000)\n",
    "\n",
    "# model_jt = RaftMLP_jt(\n",
    "#         layers = [\n",
    "#             {\"depth\": 12,\n",
    "#             \"dim\": 768,\n",
    "#             \"patch_size\": 16,\n",
    "#             \"raft_size\": 4}\n",
    "#         ],\n",
    "#         gap = True\n",
    "#     )\n",
    "\n",
    "# model_jt = SparseMLP_jt(\n",
    "#         image_size=224,\n",
    "#         patch_size=4,\n",
    "#         in_channels=3,\n",
    "#         num_classes=1000,\n",
    "#         d_model=96,\n",
    "#         depth=[2,10,24,2],\n",
    "#         expansion_factor = 2,\n",
    "#         patcher_norm= True\n",
    "#     )\n",
    "\n",
    "# model_jt = HireMLP_jt(\n",
    "#         patch_size=4,\n",
    "#         in_channels=3,\n",
    "#         num_classes=1000,\n",
    "#         d_model=[64, 128, 320, 512],\n",
    "#         h = [4,3,3,2],\n",
    "#         w = [4,3,3,2],\n",
    "#         cross_region_step = [2,2,1,1],\n",
    "#         cross_region_interval = 2,\n",
    "#         depth=[4,6,24,3],\n",
    "#         expansion_factor = 2,\n",
    "#         patcher_norm = True,\n",
    "#         padding_type = 'circular',\n",
    "#     )\n",
    "\n",
    "#model_jt = CycleMLP_B2_jt()\n",
    "\n",
    "# model_jt = AS_MLP_jt()\n",
    "# model_jt = S2MLPv1_wide()\n",
    "# model_jt = SwinMLP_jt()\n",
    "\n",
    "# model_jt = WaveMLP_jt('M')\n",
    "# model_jt = MS_MLP_jt()\n",
    "# model_jt = MorphMLP_jt('B')\n",
    "# model_jt = DynaMixer_jt('M')\n",
    "model_jt = Sequencer2D_jt('M')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "# model_pt = gMLP_pt(\n",
    "#     image_size=(224,224),\n",
    "#     patch_size=16,\n",
    "#     in_channels=3,\n",
    "#     num_classes=1000,\n",
    "#     d_model=256,\n",
    "#     d_ffn=1536,\n",
    "#     depth=30\n",
    "# )\n",
    "\n",
    "# model_pt = ConvMLP_s_pt(pretrained = True, num_classes = 1000)\n",
    "\n",
    "# model_pt = RaftMLP_pt(\n",
    "#         layers = [\n",
    "#             {\"depth\": 12,\n",
    "#             \"dim\": 768,\n",
    "#             \"patch_size\": 16,\n",
    "#             \"raft_size\": 4}\n",
    "#         ],\n",
    "#         gap = True\n",
    "#     )\n",
    "\n",
    "# model_pt = SparseMLP_pt(\n",
    "#         image_size=224,\n",
    "#         patch_size=4,\n",
    "#         in_channels=3,\n",
    "#         num_classes=1000,\n",
    "#         d_model=96,\n",
    "#         depth=[2,10,24,2],\n",
    "#         expansion_factor = 2,\n",
    "#         patcher_norm= True\n",
    "#     )\n",
    "\n",
    "# model_pt = HireMLP_pt(\n",
    "#         patch_size=4,\n",
    "#         in_channels=3,\n",
    "#         num_classes=1000,\n",
    "#         d_model=[64, 128, 320, 512],\n",
    "#         h = [4,3,3,2],\n",
    "#         w = [4,3,3,2],\n",
    "#         cross_region_step = [2,2,1,1],\n",
    "#         cross_region_interval = 2,\n",
    "#         depth=[4,6,24,3],\n",
    "#         expansion_factor = 2,\n",
    "#         patcher_norm = True,\n",
    "#     )\n",
    "\n",
    "# model_pt = GFNet_pt()\n",
    "\n",
    "# model_pt = CycleMLP_B2_pt()\n",
    "\n",
    "\n",
    "# model_pt = AS_MLP_pt()\n",
    "# model_pt = WaveMLP_S()\n",
    "# model_pt = SwinMLP_pt()\n",
    "\n",
    "# model_pt = WaveMLP_pt('M')\n",
    "# model_pt = MS_MLP_pt()\n",
    "# model_pt = MorphMLP_pt('B')\n",
    "# model_pt = DynaMixer_pt('M')\n",
    "model_pt = Sequencer2D_pt('M')\n",
    "\n",
    "# model_pt = ConvMixer_pt(\n",
    "#     dim = 1568,\n",
    "#     depth = 20\n",
    "# )\n",
    "\n",
    "# model_pt = ResMLP_pt(\n",
    "#     in_channels = 3,\n",
    "#     image_size = (224,112),\n",
    "#     patch_size = 16,\n",
    "#     d_model = 384,\n",
    "#     depth = 12,\n",
    "#     num_classes = 1000, \n",
    "#     expansion_factor = 4\n",
    "# )\n",
    "\n",
    "# model_pt = MLPMixer_pt(\n",
    "#     image_size=(224,112),\n",
    "#     patch_size=16,\n",
    "#     in_channels=3,\n",
    "#     num_classes=1000,\n",
    "#     d_model=256,\n",
    "#     depth=12,\n",
    "# )\n",
    "\n",
    "# model_pt = ViP_pt(\n",
    "#     image_size=(224, 112),\n",
    "#     patch_size=(16, 8),\n",
    "#     in_channels=3,\n",
    "#     num_classes=1000,\n",
    "#     d_model=256,\n",
    "#     depth=30,\n",
    "#     segments = 16,\n",
    "#     weighted = False\n",
    "# )\n",
    "\n",
    "# model_pt = S2MLPv2_pt(\n",
    "#     in_channels = 3,\n",
    "#     image_size = (224,224),\n",
    "#     patch_size = [(7,7), (2,2)],\n",
    "#     d_model = [192, 384],\n",
    "#     depth = [4, 14],\n",
    "#     num_classes = 1000, \n",
    "#     expansion_factor = [3, 3]\n",
    "# )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[8,1000,]\n",
      "19,959,292 total parameters.\n",
      "19,959,292 training parameters.\n"
     ]
    }
   ],
   "source": [
    "images = jt.randn(8, 3, 224, 224)\n",
    "\n",
    "with jt.no_grad():\n",
    "    output = model_jt(images)\n",
    "print(output.shape)\n",
    "\n",
    "total_params = sum(p.numel() for p in model_jt.parameters())\n",
    "print(f'{total_params:,} total parameters.')\n",
    "total_trainable_params = sum(\n",
    "    p.numel() for p in model_jt.parameters() if p.requires_grad)\n",
    "print(f'{total_trainable_params:,} training parameters.')\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "torch.Size([8, 1000])\n",
      "19,959,292 total parameters.\n",
      "19,959,292 training parameters.\n"
     ]
    }
   ],
   "source": [
    "images = torch.randn(8, 3, 224, 224)\n",
    "\n",
    "# for AS-MLP\n",
    "model_pt = model_pt.cuda()\n",
    "images = images.cuda()\n",
    "\n",
    "with torch.no_grad():\n",
    "    output = model_pt(images)\n",
    "print(output.shape)\n",
    "\n",
    "total_params = sum(p.numel() for p in model_pt.parameters())\n",
    "print(f'{total_params:,} total parameters.')\n",
    "total_trainable_params = sum(\n",
    "    p.numel() for p in model_pt.parameters() if p.requires_grad)\n",
    "print(f'{total_trainable_params:,} training parameters.')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from torchsummary import summary\n",
    "# summary(model_pt, (3, 224, 224), device=\"cpu\")\n",
    "\n",
    "total_params = sum(p.numel() for p in model_pt.parameters())\n",
    "print(f'{total_params:,} total parameters.')\n",
    "total_trainable_params = sum(\n",
    "    p.numel() for p in model_pt.parameters() if p.requires_grad)\n",
    "print(f'{total_trainable_params:,} training parameters.')\n",
    "\n",
    "# from torchstat import stat\n",
    "# stat(model_pt, (3, 224, 224))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 时间性能测试方法"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "jt.flags.use_cuda = 1\n",
    "\n",
    "bs = 32\n",
    "test_img = np.random.random((bs,3,224,224)).astype('float32')\n",
    "\n",
    "# 定义 pytorch & jittor 输入矩阵\n",
    "pytorch_test_img = torch.Tensor(test_img).cuda()\n",
    "jittor_test_img = jt.array(test_img)\n",
    "\n",
    "# 跑turns次前向求平均值\n",
    "turns = 100\n",
    "\n",
    "# 定义 pytorch & jittor 的xxx模型，如vgg\n",
    "pytorch_model = ConvMixer_pt(\n",
    "                    dim = 1568,\n",
    "                    depth = 20\n",
    "                ).cuda()\n",
    "jittor_model = ConvMixer_jt(\n",
    "                    dim = 1568,\n",
    "                    depth = 20\n",
    "                )\n",
    "\n",
    "# 把模型都设置为eval来防止dropout层对输出结果的随机影响\n",
    "pytorch_model.eval()\n",
    "jittor_model.eval()\n",
    "\n",
    "# jittor加载pytorch的初始化参数来保证参数完全相同\n",
    "jittor_model.load_parameters(pytorch_model.state_dict())\n",
    "\n",
    "key = \"ConvMixer\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# 测试Pytorch一次前向传播的平均用时\n",
    "for i in range(10):\n",
    "    pytorch_result = pytorch_model(pytorch_test_img) # Pytorch热身\n",
    "torch.cuda.synchronize()\n",
    "sta = time.time()\n",
    "for i in range(turns):\n",
    "    pytorch_result = pytorch_model(pytorch_test_img)\n",
    "torch.cuda.synchronize() # 只有运行了torch.cuda.synchronize()才会真正地运行，时间才是有效的，因此执行forward前后都要执行这句话\n",
    "end = time.time()\n",
    "tc_time = round((end - sta) / turns, 5) # 执行turns次的平均时间，输出时保留5位小数\n",
    "tc_fps = round(bs * turns / (end - sta),0) # 计算FPS\n",
    "print(f\"- Pytorch {key} forward average time cost: {tc_time}, Batch Size: {bs}, FPS: {tc_fps}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# 测试Jittor一次前向传播的平均用时\n",
    "for i in range(10):\n",
    "    jittor_result = jittor_model(jittor_test_img) # Jittor热身\n",
    "    jittor_result.sync()\n",
    "jt.sync_all(True)\n",
    "# sync_all(true)是把计算图发射到计算设备上，并且同步。只有运行了jt.sync_all(True)才会真正地运行，时间才是有效的，因此执行forward前后都要执行这句话\n",
    "sta = time.time()\n",
    "for i in range(turns):\n",
    "    jittor_result = jittor_model(jittor_test_img)\n",
    "    jittor_result.sync() # sync是把计算图发送到计算设备上\n",
    "jt.sync_all(True)\n",
    "end = time.time()\n",
    "jt_time = round((time.time() - sta) / turns, 5) # 执行turns次的平均时间，输出时保留5位小数\n",
    "jt_fps = round(bs * turns / (end - sta),0) # 计算FPS\n",
    "print(f\"- Jittor {key} forward average time cost: {jt_time}, Batch Size: {bs}, FPS: {jt_fps}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "threshold = 1e-3\n",
    "# 计算 pytorch & jittor 前向结果相对误差. 如果误差小于threshold，则测试通过.\n",
    "x = pytorch_result.detach().cpu().numpy() + 1\n",
    "y = jittor_result.data + 1\n",
    "relative_error = abs(x - y) / abs(y)\n",
    "diff = relative_error.mean()\n",
    "assert diff < threshold, f\"[*] {key} forward fails..., Relative Error: {diff}\"\n",
    "print(f\"[*] {key} forword passes with Relative Error {diff}\")"
   ]
  }
 ],
 "metadata": {
  "interpreter": {
   "hash": "40d3a090f54c6569ab1632332b64b2c03c39dcf918b08424e98f38b5ae0af88f"
  },
  "kernelspec": {
   "display_name": "Python 3.7.6 64-bit",
   "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.6.7"
  },
  "orig_nbformat": 4
 },
 "nbformat": 4,
 "nbformat_minor": 2
}