docs/tutorials/pytorch/text-classification/pruning.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Introduction"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This tutorial demonstrates how to use the pruning approach (Magnitude)based on [Intel® Neural Compressor](https://github.com/intel/neural-compressor) and benchmark of the text-classification models."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Prerequisite"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Install packages"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "* Follow [installation](https://github.com/intel-innersource/frameworks.ai.nlp-toolkit.intel-nlp-toolkit#installation) to install **nlp-toolkit**. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# install model dependency\n",
    "!pip install datasets>=1.8.0 torch>=1.10.0 transformers>=4.12.0 wandb"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Import packages"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import datasets\n",
    "import logging\n",
    "import os\n",
    "import random\n",
    "import numpy as np\n",
    "import sys\n",
    "import transformers\n",
    "from dataclasses import dataclass, field\n",
    "from datasets import load_dataset, load_metric\n",
    "from intel_extension_for_transformers.transformers import (\n",
    "    metrics,\n",
    "    OptimizedModel\n",
    ")\n",
    "from neural_compressor.config import WeightPruningConfig\n",
    "from intel_extension_for_transformers.transformers.trainer import NLPTrainer\n",
    "from transformers import (\n",
    "    AutoConfig,\n",
    "    AutoModelForSequenceClassification,\n",
    "    AutoTokenizer,\n",
    "    DataCollatorWithPadding,\n",
    "    EvalPrediction,\n",
    "    HfArgumentParser,\n",
    "    PretrainedConfig,\n",
    "    TrainingArguments,\n",
    "    default_data_collator,\n",
    "    set_seed,\n",
    ")\n",
    "from transformers.trainer_utils import get_last_checkpoint\n",
    "from transformers.utils import check_min_version\n",
    "from transformers.utils.fx import symbolic_trace\n",
    "from typing import Optional\n",
    "\n",
    "\n",
    "os.environ[\"CUDA_VISIBLE_DEVICES\"] = \"\"\n",
    "os.environ[\"WANDB_DISABLED\"] = \"true\"\n",
    "\n",
    "# Will error if the minimal version of Transformers is not installed. Remove at your own risks.\n",
    "check_min_version(\"4.12.0\")\n",
    "\n",
    "\n",
    "task_to_keys = {\n",
    "    \"cola\": (\"sentence\", None),\n",
    "    \"mnli\": (\"premise\", \"hypothesis\"),\n",
    "    \"mrpc\": (\"sentence1\", \"sentence2\"),\n",
    "    \"qnli\": (\"question\", \"sentence\"),\n",
    "    \"qqp\": (\"question1\", \"question2\"),\n",
    "    \"rte\": (\"sentence1\", \"sentence2\"),\n",
    "    \"sst2\": (\"sentence\", None),\n",
    "    \"stsb\": (\"sentence1\", \"sentence2\"),\n",
    "    \"wnli\": (\"sentence1\", \"sentence2\"),\n",
    "}\n",
    "\n",
    "logger = logging.getLogger(__name__)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Define arguments"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# ========== Define arguments =========\n",
    "@dataclass\n",
    "class DataTrainingArguments:\n",
    "    \"\"\"\n",
    "    Arguments pertaining to what data we are going to input our model for training and eval.\n",
    "    Using `HfArgumentParser` we can turn this class\n",
    "    into argparse arguments to be able to specify them on\n",
    "    the command line.\n",
    "    \"\"\"\n",
    "\n",
    "    task_name: Optional[str] = field(\n",
    "        default=None,\n",
    "        metadata={\"help\": \"The name of the task to train on: \" + \", \".join(task_to_keys.keys())},\n",
    "    )\n",
    "    dataset_name: Optional[str] = field(\n",
    "        default=None, metadata={\"help\": \"The name of the dataset to use (via the datasets library).\"}\n",
    "    )\n",
    "    dataset_config_name: Optional[str] = field(\n",
    "        default=None, metadata={\"help\": \"The configuration name of the dataset to use (via the datasets library).\"}\n",
    "    )\n",
    "    max_seq_length: int = field(\n",
    "        default=128,\n",
    "        metadata={\n",
    "            \"help\": \"The maximum total input sequence length after tokenization. Sequences longer \"\n",
    "            \"than this will be truncated, sequences shorter will be padded.\"\n",
    "        },\n",
    "    )\n",
    "    overwrite_cache: bool = field(\n",
    "        default=False, metadata={\"help\": \"Overwrite the cached preprocessed datasets or not.\"}\n",
    "    )\n",
    "    pad_to_max_length: bool = field(\n",
    "        default=True,\n",
    "        metadata={\n",
    "            \"help\": \"Whether to pad all samples to `max_seq_length`. \"\n",
    "            \"If False, will pad the samples dynamically when batching to the maximum length in the batch.\"\n",
    "        },\n",
    "    )\n",
    "    max_train_samples: Optional[int] = field(\n",
    "        default=None,\n",
    "        metadata={\n",
    "            \"help\": \"For debugging purposes or quicker training, truncate the number of training examples to this \"\n",
    "            \"value if set.\"\n",
    "        },\n",
    "    )\n",
    "    max_eval_samples: Optional[int] = field(\n",
    "        default=None,\n",
    "        metadata={\n",
    "            \"help\": \"For debugging purposes or quicker training, truncate the number of evaluation examples to this \"\n",
    "            \"value if set.\"\n",
    "        },\n",
    "    )\n",
    "    max_predict_samples: Optional[int] = field(\n",
    "        default=None,\n",
    "        metadata={\n",
    "            \"help\": \"For debugging purposes or quicker training, truncate the number of prediction examples to this \"\n",
    "            \"value if set.\"\n",
    "        },\n",
    "    )\n",
    "    train_file: Optional[str] = field(\n",
    "        default=None, metadata={\"help\": \"A csv or a json file containing the training data.\"}\n",
    "    )\n",
    "    validation_file: Optional[str] = field(\n",
    "        default=None, metadata={\"help\": \"A csv or a json file containing the validation data.\"}\n",
    "    )\n",
    "\n",
    "    def __post_init__(self):\n",
    "        if self.task_name is not None:\n",
    "            self.task_name = self.task_name.lower()\n",
    "            if self.task_name not in task_to_keys.keys():\n",
    "                raise ValueError(\"Unknown task, you should pick one in \" + \",\".join(task_to_keys.keys()))\n",
    "        elif self.dataset_name is not None:\n",
    "            pass\n",
    "        elif self.train_file is None or self.validation_file is None:\n",
    "            raise ValueError(\"Need either a GLUE task, a training/validation file or a dataset name.\")\n",
    "        else:\n",
    "            train_extension = self.train_file.split(\".\")[-1]\n",
    "            assert train_extension in [\"csv\", \"json\"], \"`train_file` should be a csv or a json file.\"\n",
    "            validation_extension = self.validation_file.split(\".\")[-1]\n",
    "            assert (\n",
    "                validation_extension == train_extension\n",
    "            ), \"`validation_file` should have the same extension (csv or json) as `train_file`.\"\n",
    "\n",
    "\n",
    "@dataclass\n",
    "class ModelArguments:\n",
    "    \"\"\"\n",
    "    Arguments pertaining to which model/config/tokenizer we are going to fine-tune from.\n",
    "    \"\"\"\n",
    "\n",
    "    model_name_or_path: str = field(\n",
    "        metadata={\"help\": \"Path to pretrained model or model identifier from huggingface.co/models\"}\n",
    "    )\n",
    "\n",
    "\n",
    "@dataclass\n",
    "class OptimizationArguments:\n",
    "    \"\"\"\n",
    "    Arguments pertaining to what type of optimization we are going to apply on the model.\n",
    "    \"\"\"\n",
    "\n",
    "    prune: bool = field(\n",
    "        default=False,\n",
    "        metadata={\"help\": \"Whether or not to apply prune.\"},\n",
    "    )\n",
    "    pruning_approach: Optional[str] = field(\n",
    "        default=\"BasicMagnitude\",\n",
    "        metadata={\"help\": \"Pruning approach. Supported approach is basic_magnite.\"},\n",
    "    )\n",
    "    target_sparsity_ratio: Optional[float] = field(\n",
    "        default=None,\n",
    "        metadata={\"help\": \"Targeted sparsity when pruning the model.\"},\n",
    "    )\n",
    "    metric_name: Optional[str] = field(\n",
    "        default=None,\n",
    "        metadata={\"help\": \"Metric used for the tuning strategy.\"},\n",
    "    )\n",
    "    tolerance_mode: Optional[str] = field(\n",
    "        default=\"absolute\",\n",
    "        metadata={\"help\": \"Metric tolerance model, expected to be relative or absolute.\"},\n",
    "    )\n",
    "    perf_tol: Optional[float] = field(\n",
    "        default=0.02,\n",
    "        metadata={\"help\": \"Performance tolerance when optimizing the model.\"},\n",
    "    )\n",
    "    benchmark: bool = field(\n",
    "        default=False,\n",
    "        metadata={\"help\": \"run benchmark.\"})\n",
    "    accuracy_only: bool = field(\n",
    "        default=False,\n",
    "        metadata={\"help\":\"Whether to only test accuracy for model tuned by Neural Compressor.\"})"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "model_args = ModelArguments(\n",
    "    model_name_or_path=\"textattack/bert-base-uncased-MRPC\",\n",
    ")\n",
    "data_args = DataTrainingArguments(\n",
    "    task_name=\"mrpc\",\n",
    "    max_seq_length=128,\n",
    "    overwrite_cache=True\n",
    ")\n",
    "training_args = TrainingArguments(\n",
    "    output_dir=\"./saved_result\",\n",
    "    do_eval=True,\n",
    "    do_train=True,\n",
    "    no_cuda=True,\n",
    "    overwrite_output_dir=True,\n",
    "    per_device_train_batch_size=8,\n",
    ")\n",
    "optim_args = OptimizationArguments(\n",
    "    tune=True,\n",
    "    quantization_approach=\"static\"\n",
    ")\n",
    "log_level = training_args.get_process_log_level()\n",
    "logger.setLevel(log_level)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Download dataset from the hub"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# download the dataset.\n",
    "raw_datasets = load_dataset(\"glue\", data_args.task_name)\n",
    "# Labels\n",
    "label_list = raw_datasets[\"train\"].features[\"label\"].names\n",
    "num_labels = len(label_list)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Download fp32 model from the hub"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Load pretrained model and tokenizer\n",
    "config = AutoConfig.from_pretrained(\n",
    "    model_args.model_name_or_path,\n",
    "    num_labels=num_labels,\n",
    "    finetuning_task=data_args.task_name,\n",
    "    revision=\"main\"\n",
    ")\n",
    "tokenizer = AutoTokenizer.from_pretrained(\n",
    "    model_args.model_name_or_path,\n",
    "    use_fast=True,\n",
    "    revision=\"main\"\n",
    ")\n",
    "model = AutoModelForSequenceClassification.from_pretrained(\n",
    "    model_args.model_name_or_path,\n",
    "    from_tf=bool(\".ckpt\" in model_args.model_name_or_path),\n",
    "    config=config,\n",
    "    revision=\"main\"\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Preprocessing the dataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Preprocessing the raw_datasets\n",
    "sentence1_key, sentence2_key = task_to_keys[data_args.task_name]\n",
    "# Padding strategy\n",
    "padding = False\n",
    "# Some models have set the order of the labels to use, so let's make sure we do use it.\n",
    "label_to_id = None\n",
    "if model.config.label2id != PretrainedConfig(num_labels=num_labels).label2id:\n",
    "    # Some have all caps in their config, some don't.\n",
    "    label_name_to_id = {k.lower(): v for k, v in model.config.label2id.items()}\n",
    "    if list(sorted(label_name_to_id.keys())) == list(sorted(label_list)):\n",
    "        label_to_id = {i: int(label_name_to_id[label_list[i]]) for i in range(num_labels)}\n",
    "    else:\n",
    "        logger.warning(\n",
    "            f\"Your model seems to have been trained with labels, but they don't match the dataset: \"\n",
    "            f\"model labels: {list(sorted(label_name_to_id.keys()))}, dataset labels: {list(sorted(label_list))}.\\n\"\n",
    "            f\"Ignoring the model labels as a result.\"\n",
    "        )\n",
    "if label_to_id is not None:\n",
    "    model.config.label2id = label_to_id\n",
    "    model.config.id2label = {id: label for label, id in config.label2id.items()}\n",
    "max_seq_length = min(data_args.max_seq_length, tokenizer.model_max_length)\n",
    "\n",
    "def preprocess_function(examples):\n",
    "    # Tokenize the texts\n",
    "    args = (\n",
    "        (examples[sentence1_key],) if sentence2_key is None else (examples[sentence1_key], examples[sentence2_key])\n",
    "    )\n",
    "    result = tokenizer(*args, padding=padding, max_length=max_seq_length, truncation=True)\n",
    "\n",
    "    # Map labels to IDs (not necessary for GLUE tasks)\n",
    "    if label_to_id is not None and \"label\" in examples:\n",
    "        result[\"label\"] = [(label_to_id[l] if l != -1 else -1) for l in examples[\"label\"]]\n",
    "    return result\n",
    "\n",
    "with training_args.main_process_first(desc=\"dataset map pre-processing\"):\n",
    "    raw_datasets = raw_datasets.map(\n",
    "        preprocess_function, batched=True, load_from_cache_file=not data_args.overwrite_cache\n",
    "    )\n",
    "\n",
    "if training_args.do_train:\n",
    "    if \"train\" not in raw_datasets:\n",
    "        raise ValueError(\"--do_train requires a train dataset\")\n",
    "    train_dataset = raw_datasets[\"train\"]\n",
    "\n",
    "if training_args.do_eval:\n",
    "    if \"validation\" not in raw_datasets and \"validation_matched\" not in raw_datasets:\n",
    "        raise ValueError(\"--do_eval requires a validation dataset\")\n",
    "    eval_dataset = raw_datasets[\"validation_matched\" if data_args.task_name == \"mnli\" else \"validation\"]\n",
    "\n",
    "# Log a few random samples from the training set:\n",
    "if training_args.do_train:\n",
    "    for index in random.sample(range(len(train_dataset)), 3):\n",
    "        logger.info(f\"Sample {index} of the training set: {train_dataset[index]}.\")\n",
    "\n",
    "# Get the metric function\n",
    "metric = load_metric(\"glue\", data_args.task_name)\n",
    "\n",
    "metric_name = \"eval_accuracy\"\n",
    "\n",
    "# You can define your custom compute_metrics function. It takes an `EvalPrediction` object (a namedtuple with a\n",
    "# predictions and label_ids field) and has to return a dictionary string to float.\n",
    "def compute_metrics(p: EvalPrediction):\n",
    "    preds = p.predictions[0] if isinstance(p.predictions, tuple) else p.predictions\n",
    "    preds =  np.argmax(preds, axis=1)\n",
    "    if data_args.task_name is not None:\n",
    "        result = metric.compute(predictions=preds, references=p.label_ids)\n",
    "        if len(result) > 1:\n",
    "            result[\"combined_score\"] = np.mean(list(result.values())).item()\n",
    "        return result\n",
    "    else:\n",
    "        return {\"accuracy\": (preds == p.label_ids).astype(np.float32).mean().item()}\n",
    "\n",
    "# Data collator will default to DataCollatorWithPadding, so we change it if we already did the padding.\n",
    "data_collator = None"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Pruning & Benchmark"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Pruning"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Set seed before initializing model.\n",
    "set_seed(training_args.seed)\n",
    "# Initialize our Trainer\n",
    "trainer = NLPTrainer(\n",
    "    model=model,\n",
    "    args=training_args,\n",
    "    train_dataset=train_dataset if training_args.do_train else None,\n",
    "    eval_dataset=eval_dataset if training_args.do_eval else None,\n",
    "    compute_metrics=compute_metrics,\n",
    "    tokenizer=tokenizer,\n",
    "    data_collator=data_collator,\n",
    ")\n",
    "\n",
    "tune_metric = metrics.Metric(name=metric_name)\n",
    "prune_type = 'magnitude' if optim_args.pruning_approach else optim_args.pruning_approach\n",
    "target_sparsity_ratio = optim_args.target_sparsity_ratio \\\n",
    "    if optim_args.target_sparsity_ratio else None\n",
    "trainer.metrics = tune_metric\n",
    "pruning_conf = WeightPruningConfig([{\"start_step\": 0, \"end_step\": 2}],\n",
    "                                    target_sparsity=target_sparsity_ratio,\n",
    "                                    pruning_scope=\"local\",\n",
    "                                    pruning_type=prune_type)\n",
    "\n",
    "model = trainer.prune(pruning_config=pruning_conf)\n",
    "trainer.save_model(training_args.output_dir)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Run Benchmark after Pruning"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Load the model obtained after Intel Neural Compressor (INC) quantization\n",
    "model = OptimizedModel.from_pretrained(\n",
    "    training_args.output_dir,\n",
    ")\n",
    "model.eval()\n",
    "trainer.model = model\n",
    "results = trainer.evaluate()\n",
    "logger.info(\"metrics keys: {}\".format(results.keys()))\n",
    "bert_task_acc_keys = ['eval_f1', 'eval_accuracy', 'eval_matthews_correlation',\n",
    "                        'eval_pearson', 'eval_mcc', 'eval_spearmanr']\n",
    "for key in bert_task_acc_keys:\n",
    "    if key in results.keys():\n",
    "        throughput = results.get(\"eval_samples_per_second\")\n",
    "        print('Batch size = ', training_args.per_device_eval_batch_size)\n",
    "        print(\"Finally Eval {} Accuracy: {}\".format(key, results[key]))\n",
    "        print(\"Latency: {:.5f} ms\".format(1000 / throughput))\n",
    "        print(\"Throughput: {:.5f} samples/sec\".format(throughput))"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3.10.6 64-bit",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "name": "python",
   "version": "3.10.6"
  },
  "orig_nbformat": 4,
  "vscode": {
   "interpreter": {
    "hash": "a3ed54c68abdb79eabea0140062ffa976ea4d8132b937aa83ca919a8d862edf2"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}