image_task/example/test_speed.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### We provide this demo to show the speeding up of R-MeeTo model\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "source": [
    "import sys\n",
    "import os\n",
    "import timm\n",
    "import torch\n",
    "from timm.models import create_model\n",
    "\n",
    "sys.path.append(os.path.abspath(os.path.join(os.getcwd(), '..')))\n",
    "from R_MeeTo import RMeeTo_Mamba,RMeeTo_Mamba_Video,RMeeToVisionTransformer\n",
    "from utils.calc_flops import throughput,get_flops\n"
   ],
   "outputs": []
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "source": [
    "#choose the model\n",
    "model_merge = 'RMeeTo_small'\n",
    "model_origin = model_merge+'_teacher'\n",
    "\n",
    "model_pth = './pretrained'"
   ],
   "outputs": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "source": [
    "#set merge parameters\n",
    "merge_para = {\n",
    "        'num_prune': 11,\n",
    "        'distance': 'cosine',\n",
    "        'metric': 'X',\n",
    "        'merge_interval': 2,\n",
    "        'if_prune': False,\n",
    "        'if_order': True,\n",
    "        'class_token': True,\n",
    "        'if_merge': True,\n",
    "        'if_merge_odd': False,\n",
    "        'merge_mode':'sum',\n",
    "        'choose':'max',\n",
    "        'compare': 11,\n",
    "    }\n",
    "\n",
    "print(\"Merge para:\",merge_para)"
   ],
   "outputs": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "source": [
    "model = create_model(\n",
    "    model_merge,\n",
    "    pretrained=True,\n",
    "    model_pth=model_pth,\n",
    "    merge_para=merge_para\n",
    ")\n",
    "\n",
    "origin_model = create_model(\n",
    "    model_origin,\n",
    "    pretrained=True,\n",
    "    model_pth=model_pth\n",
    ")\n",
    "\n",
    "device = 'cuda'\n",
    "model.to(device)\n",
    "origin_model.to(device)"
   ],
   "outputs": []
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "source": [
    "#get flops\n",
    "print('Get FLOPs G) for origin model')\n",
    "get_flops(model=origin_model,show_detail=False)\n",
    "print('Get FLOPs G) for R-MeeTo model')\n",
    "get_flops(model=model,show_detail=False)\n"
   ],
   "outputs": []
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "source": [
    "#get throughput\n",
    "images = torch.randn(128, 3,224, 224).to(device)\n",
    "print('Get throughput for origin model')\n",
    "throughput(images,model=origin_model)\n",
    "print('Get throughput for R-MeeTo model')\n",
    "throughput(images,model=model)"
   ],
   "outputs": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "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.10.13"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}