From f6043b2d06858c9a2c408f6102fcfda0bf0b5336 Mon Sep 17 00:00:00 2001 From: "D.C." <71894595+semi-dlc@users.noreply.github.com> Date: Thu, 30 Jul 2026 19:04:18 +0200 Subject: [PATCH 1/3] fix: initialize quantizer parameters directly on cuda device fix: initialize quantizer parameters directly on cuda device squashed and added pre-commit changes Fixes a bug in the hyperparameter optimization loop with the jetset-classifier, which moved parameters of the Quantizer class from cuda to cpu at reinitialization. --- src/pquant/core/torch/quantizer.py | 10 +-- tests/test_reinitialize_quantizers.py | 110 ++++++++++++++++++++++++++ 2 files changed, 115 insertions(+), 5 deletions(-) create mode 100644 tests/test_reinitialize_quantizers.py diff --git a/src/pquant/core/torch/quantizer.py b/src/pquant/core/torch/quantizer.py index b664547..51abaf5 100644 --- a/src/pquant/core/torch/quantizer.py +++ b/src/pquant/core/torch/quantizer.py @@ -23,7 +23,7 @@ def __init__( shape=None, ): super().__init__() - + device = "cuda" if torch.cuda.is_available() else "cpu" self.overflow = overflow self.round_mode = round_mode self.use_hgq = is_heterogeneous @@ -32,10 +32,10 @@ def __init__( self.granularity = QuantizationGranularity(granularity).value if not self.use_hgq: param_shape = () if is_data else self.compute_weight_param_shape(shape) - self.k = torch.nn.Parameter(torch.full(param_shape, float(k)), requires_grad=False) - self.i = torch.nn.Parameter(torch.full(param_shape, float(i)), requires_grad=False) - self.f = torch.nn.Parameter(torch.full(param_shape, float(f)), requires_grad=False) - self.b = torch.nn.Parameter(torch.full(param_shape, float(i + k + f)), requires_grad=False) + self.k = torch.nn.Parameter(torch.full(param_shape, float(k), device=device), requires_grad=False) + self.i = torch.nn.Parameter(torch.full(param_shape, float(i), device=device), requires_grad=False) + self.f = torch.nn.Parameter(torch.full(param_shape, float(f), device=device), requires_grad=False) + self.b = torch.nn.Parameter(torch.full(param_shape, float(i + k + f), device=device), requires_grad=False) self.quantizer = create_quantizer( k, i, diff --git a/tests/test_reinitialize_quantizers.py b/tests/test_reinitialize_quantizers.py new file mode 100644 index 0000000..2e56dc0 --- /dev/null +++ b/tests/test_reinitialize_quantizers.py @@ -0,0 +1,110 @@ +# requires CUDA as this test mainly tests whether moving Quantizers to CUDA devices work +import pytest +import torch + +from pquant.core.torch.quantizer import Quantizer + +CUDA_AVAILABLE = torch.cuda.is_available() + + +def make_default_quantizer(**overrides): + """ + Default quantizer used similiarly in training loop. + """ + kwargs = dict( + k=1, + i=4, + f=7, + overflow="SAT", + round_mode="RND", + is_heterogeneous=False, + is_data=False, + granularity="per_channel", + hgq_gamma=0.0003, + place="datalane", + dynamic_data=True, + ) + kwargs.update(overrides) + return Quantizer(**kwargs) + + +def assert_all_params_on(module: Quantizer, device: str): + """ + Assert that all registered quantizing variables (b, f, k, i) of quantizer are on device. + """ + for name, param in module.named_parameters(): + if name in ("b", "f", "k", "i"): + assert param.device.type == device, f"Parameter '{name}' is on {param.device}, expected {device}" + + +class TestQuantizerDevices: + def setup_method(self): + self.device = "cuda" if CUDA_AVAILABLE else "cpu" + + def test_initial_construction_device_choice(self): + """ + Tests that quantizing variables (b, f, k, i) are on the right device + """ + q = make_default_quantizer() + expected = "cuda" if CUDA_AVAILABLE else "cpu" + assert_all_params_on(q, expected) + + @pytest.mark.skipif(not CUDA_AVAILABLE, reason="requires CUDA to test CUDA placement") + def test_reconstruction(self): + """ + Tests reinitialization of quantizer + """ + q = make_default_quantizer() + for _ in range(3): + q = make_default_quantizer() + assert_all_params_on(q, "cuda") + + @pytest.mark.skipif(not CUDA_AVAILABLE, reason="requires CUDA to test CUDA placement") + def test_set_quantization_bits_preserves_cuda(self): + """ + Tests if calling set_quantization_bits changes device + """ + q = make_default_quantizer() + # Ensure initial parameters are all on CUDA + assert_all_params_on(q, "cuda") + + # Call the real implementation + q.set_quantization_bits(i=2, f=9) + + # After the call, all registered parameters should still be on CUDA + assert_all_params_on(q, "cuda") + + @pytest.mark.skipif(not CUDA_AVAILABLE, reason="requires CUDA to test CUDA placement") + def test_apply_final_compression(self): + """ + Tests if calling apply_final_compression changes device + """ + q = make_default_quantizer() + q.to(self.device) + assert_all_params_on(q, "cuda") + + # This will call get_quantization_bits and then reassign i, f, b, + # and final_compression_done.data. + q.apply_final_compression() + + assert_all_params_on(q, "cuda") + + @pytest.mark.skipif(not CUDA_AVAILABLE, reason="requires CUDA to test CUDA placement") + def test_quantized_forward_pass(self): + """ + Tests whether the output of a forward pass through a quantized network is still on device. + If any of the bit parameters ended up on CPU, this should raise + a device mismatch error. + """ + q = make_default_quantizer() + q.to(self.device) + assert_all_params_on(q, "cuda") + + q.set_quantization_bits(i=3, f=5) + assert_all_params_on(q, "cuda") + + # dummy input on CUDA + x = torch.randn(4, 8, device=torch.device("cuda")) + + out = q(x) + assert out.device.type == "cuda", "Output tensor is not on CUDA as expected" From 3d6621de78d1e9ea82a83d1e8cebe22a259e5ac3 Mon Sep 17 00:00:00 2001 From: "D.C." <71894595+semi-dlc@users.noreply.github.com> Date: Fri, 31 Jul 2026 11:45:54 +0200 Subject: [PATCH 2/3] Added reinitialize_quantizers test to run_test.sh --- tests/run_tests.sh | 1 + 1 file changed, 1 insertion(+) diff --git a/tests/run_tests.sh b/tests/run_tests.sh index 34fedbc..01d585a 100755 --- a/tests/run_tests.sh +++ b/tests/run_tests.sh @@ -16,3 +16,4 @@ KERAS_BACKEND=torch pytest test_torch_alkaid_conversion.py pytest test_keras_alkaid_conversion.py KERAS_BACKEND=torch pytest test_hgq_torch.py pytest test_hgq_keras.py +KERAS_BACKEND=torch pytest test_reinitialize_quantizers.py From 8d30a640d17cb98f22324f24ab02ff3bf96d9935 Mon Sep 17 00:00:00 2001 From: "D.C." <71894595+semi-dlc@users.noreply.github.com> Date: Fri, 21 Aug 2026 15:27:06 +0200 Subject: [PATCH 3/3] Documentation of summer student project (code) --- examples/example_transformer_jet_tagger.ipynb | 2377 +++++++++++++++++ 1 file changed, 2377 insertions(+) create mode 100644 examples/example_transformer_jet_tagger.ipynb diff --git a/examples/example_transformer_jet_tagger.ipynb b/examples/example_transformer_jet_tagger.ipynb new file mode 100644 index 0000000..65825d0 --- /dev/null +++ b/examples/example_transformer_jet_tagger.ipynb @@ -0,0 +1,2377 @@ +{ + "cells": [ + { + "cell_type": "markdown", + "id": "8dc39bab", + "metadata": {}, + "source": [ + "# Jet-tagging with transformers\n", + "Notebook summarizing the most relevant parts of the openlab summer student project by Dino Cheng in 2026 at EP/SFT CERN, working on compressing a transformer-based jet tagger from heptokens with PQuantML using pruning via PDP/DST, quantization, and FITcompress.\n", + "It trains variants of a transformer-based jet tagger with various compression methods, both successful and unsuccessful.\n", + "A full report is available on Zenodo starting in (approx.) September 2026.\n", + "### Please do not execute this in the PQuantML repo! Move this into some other working directory due to import logics." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "aa437be1", + "metadata": {}, + "outputs": [], + "source": [ + "%matplotlib inline\n", + "\n", + "# pquant selects its backend from KERAS_BACKEND at import time (defaults to\n", + "# \"tensorflow\"). This project uses the torch layers (PQDense, etc.), so the\n", + "# variable must be set before any `pquant` import below.\n", + "import os\n", + "\n", + "os.environ[\"KERAS_BACKEND\"] = \"torch\"\n", + "\n", + "# generic stuff\n", + "import logging\n", + "from abc import ABC, abstractmethod\n", + "from collections.abc import Callable\n", + "from datetime import datetime\n", + "from functools import partial\n", + "import logging\n", + "import json\n", + "from pathlib import Path\n", + "from typing import Literal\n", + "\n", + "import joblib\n", + "import matplotlib\n", + "import matplotlib.pyplot as plt\n", + "import seaborn as sns\n", + "import numpy as np\n", + "import pandas as pd\n", + "\n", + "# torch\n", + "import torch\n", + "import torch as T\n", + "import torch.nn as nn\n", + "from torch.nn.functional import cross_entropy, gelu, relu\n", + "from torch.utils.data import DataLoader, Dataset, random_split\n", + "\n", + "# lightning\n", + "import lightning as L\n", + "from lightning import LightningModule\n", + "from lightning.pytorch.callbacks import LearningRateMonitor, ModelCheckpoint\n", + "from torchmetrics import AUROC, Accuracy, ConfusionMatrix, F1Score\n", + "\n", + "# onnx\n", + "import onnxruntime as ort\n", + "\n", + "# PQuantML\n", + "from pquant.quantizer import Quantizer\n", + "from pquant import get_ebops, get_layer_keep_ratio\n", + "from pquant.core.hyperparameter_optimization import dst_config, pdp_config, wanda_config, fitcompress_config\n", + "from pquant.activations import PQActivation\n", + "from pquant.core.torch import convert_to_onnx\n", + "from pquant.core.torch.layers import (\n", + " PQLayerNorm,\n", + " call_post_round_functions,\n", + " post_epoch_functions,\n", + " post_pretrain_functions,\n", + " pre_epoch_functions,\n", + " pre_finetune_functions,\n", + " save_weights_functions,\n", + ")\n", + "from pquant.core.torch.train import train_model\n", + "from pquant.core.torch.tracing import check_quantization\n", + "from pquant.layers import PQDense, PQMultiheadAttention, PQConv1d, PQConv2d, add_compression_layers, get_model_losses\n", + "\n", + "\n", + "# heptokens (Jeff's library, used for data loader functionality)\n", + "import heptokens\n", + "from heptokens.data.atlas_mappable import SingleFileMapModule\n", + "from heptokens.data.collation import preprocess_batch\n", + "from heptokens.models.token_classifier import (\n", + " Embedder,\n", + " Pooler,\n", + " SequenceEncoder,\n", + " ScheduledOptimiserMixin,\n", + ")\n", + "\n" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "6022bcf3", + "metadata": {}, + "outputs": [], + "source": [ + "def timestamp():\n", + " return datetime.now().strftime(\"%m-%d-%H-%M-%S\")" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "26a7724d", + "metadata": {}, + "outputs": [], + "source": [ + "def state_dict_export(model, path, verbose=True):\n", + " sd = model.state_dict()\n", + " if verbose:\n", + " sd_serializable = {\n", + " name: {\n", + " \"shape\": list(tensor.shape),\n", + " \"dtype\": str(tensor.dtype),\n", + " \"values\": tensor.detach().cpu().numpy().tolist(),\n", + " }\n", + " for name, tensor in sd.items()\n", + " }\n", + " with open(path, \"w\") as f:\n", + " json.dump(sd_serializable, f, indent=2)\n", + "\n", + " else:\n", + " with open(path, \"w\") as f:\n", + " for name, tensor in sd.items():\n", + " t = tensor.detach().cpu().float()\n", + " f.write(\n", + " f\"{name}\\tshape={list(t.shape)}\\tdtype={tensor.dtype}\\t\"\n", + " f\"min={t.min().item():.6f}\\tmax={t.max().item():.6f}\\t\"\n", + " f\"mean={t.mean().item():.6f}\\n\"\n", + " )\n", + " f.write(str(model.named_parameters))\n" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "7c616dfb", + "metadata": {}, + "outputs": [], + "source": [ + "PREPROCESSING_ON = False # if True: requires preprocessing scalers from heptokens" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "2ff45cd7", + "metadata": {}, + "outputs": [], + "source": [ + "# torch configs\n", + "torch.manual_seed(42)\n", + "device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n", + "print(f\"Using device: {device}\")\n", + "\n", + "torch.set_float32_matmul_precision(\"high\") # use TF32 on H100\n", + "torch.backends.cudnn.benchmark = True\n", + "torch.backends.cuda.enable_flash_sdp(True)\n", + "\n", + "dtype_real = torch.float32" + ] + }, + { + "cell_type": "markdown", + "id": "75df5730", + "metadata": {}, + "source": [ + "Some constants/names are defined in the following cell" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "fc7d4020", + "metadata": {}, + "outputs": [], + "source": [ + "TRACK_FEATURES = [\n", + " \"pt\",\n", + " \"deta\",\n", + " \"dphi\",\n", + " \"d0\",\n", + " \"d0RelativeToBeamspot\",\n", + " \"d0Uncertainty\",\n", + " \"d0RelativeToBeamspotUncertainty\",\n", + " \"z0RelativeToBeamspot\",\n", + " \"z0RelativeToBeamspotUncertainty\",\n", + " \"z0SinTheta\",\n", + " \"z0SinThetaUncertainty\",\n", + " \"lifetimeSignedD0\",\n", + " \"lifetimeSignedD0Significance\",\n", + " \"lifetimeSignedZ0SinTheta\",\n", + " \"lifetimeSignedZ0SinThetaSignificance\",\n", + " \"theta\",\n", + " \"thetaUncertainty\",\n", + " \"qOverP\",\n", + " \"qOverPUncertainty\",\n", + " \"ptfrac\",\n", + "]\n", + "\n", + "CLASS_NAMES = [\"light\", \"c\", \"b\", \"tau\"]\n", + "\n", + "JET_FEATURES = [\"pt\", \"mass\", \"eta\", \"phi\"]\n", + "\n", + "MAX_JET_PT = 7_000_000.0\n", + "MAX_CST_PT = 1_000_000.0\n", + "LABEL_KEY = \"HadronConeExclTruthLabelID\"\n", + "NUM_CSTS = 40\n", + "\n", + "LABEL_MAP = {0: 0, 4: 1, 5: 2, 15: 3} \n", + "\n", + "# in the decision tree, they are ordered like 5 4 15 0\n", + "FLAVOUR_LABELS = {\n", + " 0: \"light\",\n", + " 4: \"c\",\n", + " 5: \"b\",\n", + " 15: \"tau\",\n", + "}" + ] + }, + { + "cell_type": "markdown", + "id": "6de0e1d6", + "metadata": {}, + "source": [ + "These are plotting functions from view_model.ipynb and view_run.ipynb.\n" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "d094bfc2", + "metadata": {}, + "outputs": [], + "source": [ + "def collect_layer_bits(model):\n", + "\n", + " compressed_types = (PQConv1d, PQConv2d, PQDense)\n", + "\n", + " names = []\n", + " integer_bits = []\n", + " fractional_bits = []\n", + " integer_bias_bits = []\n", + " fractional_bias_bits = []\n", + "\n", + " for n, m in model.named_modules():\n", + " if isinstance(m, compressed_types):\n", + " k, i, f = m.get_weight_quantization_bits()\n", + " names.append(n)\n", + " integer_bits.append(float(np.mean(i.detach().cpu().numpy()) if torch.is_tensor(i) else i))\n", + " fractional_bits.append(float(np.mean(f.detach().cpu().numpy()) if torch.is_tensor(f) else f))\n", + "\n", + " if hasattr(m, \"get_bias_quantization_bits\") and m.bias is not None:\n", + " k_b, i_b, f_b = m.get_bias_quantization_bits()\n", + " integer_bias_bits.append(float(np.mean(i_b.detach().cpu().numpy()) if torch.is_tensor(i_b) else i_b))\n", + " fractional_bias_bits.append(float(np.mean(f_b.detach().cpu().numpy()) if torch.is_tensor(f_b) else f_b))\n", + " else:\n", + " integer_bias_bits.append(0.0)\n", + " fractional_bias_bits.append(0.0)\n", + "\n", + " return names, integer_bits, fractional_bits, integer_bias_bits, fractional_bias_bits\n", + "\n", + "\n", + "def plot_bit_allocation(model, use_bias=False):\n", + " names, integer_bits, fractional_bits, integer_bias_bits, fractional_bias_bits = collect_layer_bits(model)\n", + "\n", + " if not names:\n", + " raise RuntimeError(\n", + " \"No compressed layers found. Make sure add_compression_layers() and \"\n", + " \"FITCompress have been run on this model before plotting.\"\n", + " )\n", + "\n", + " ibits = np.array(integer_bias_bits if use_bias else integer_bits, dtype=float)\n", + " fbits = np.array(fractional_bias_bits if use_bias else fractional_bits, dtype=float)\n", + "\n", + " y = np.arange(len(names))\n", + " fig_h = max(4, 0.28 * len(names))\n", + " fig, ax = plt.subplots(figsize=(7, fig_h))\n", + "\n", + " ax.barh(y, -ibits, color=\"orange\", label=\"Integer bits\")\n", + " ax.barh(y, fbits, color=\"blue\", label=\"Fractional bits\")\n", + "\n", + " ax.set_yticks(y)\n", + " ax.set_yticklabels(names, fontsize=7)\n", + " ax.invert_yaxis()\n", + " ax.axvline(0, color=\"black\", linewidth=0.8)\n", + "\n", + " max_bits = max(ibits.max(), fbits.max()) + 1\n", + " ax.set_xlim(-max_bits, max_bits)\n", + " ticks = ax.get_xticks()\n", + " ax.set_xticks(ticks)\n", + " ax.set_xticklabels([str(abs(int(t))) for t in ticks])\n", + "\n", + " ax.set_xlabel(\"Bit-width\")\n", + " ax.set_title(\"Per-layer Bit Allocation (Integer vs Fractional) after FITCompress\")\n", + " ax.legend(loc=\"lower right\")\n", + " ax.grid(axis=\"x\", linestyle=\"--\", alpha=0.4)\n", + "\n", + " plt.tight_layout()\n", + " return fig, ax" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "482b4581", + "metadata": {}, + "outputs": [], + "source": [ + "def plot_remaining_weights(model):\n", + " # Plot remaining weights\n", + " names = []\n", + " remaining = []\n", + " total_w = []\n", + " nonzeros = []\n", + "\n", + " integer_bits = []\n", + " integer_bias_bits = []\n", + " fractional_bits = []\n", + " fractional_bias_bits = []\n", + "\n", + " for n, m in model.named_modules():\n", + " if isinstance(m, (torch.nn.Conv1d, torch.nn.Conv2d, torch.nn.Linear, torch.nn.ReLU)):\n", + " names.append(n.replace(\"encoder.transformer.encoder.layers.\", \"\").replace(\"encoder.transformer.\", \"\"))\n", + " nonzero = np.count_nonzero(m.weight.detach().cpu())\n", + " remaining_pct = nonzero / m.weight.numel()\n", + " remaining.append(remaining_pct)\n", + " total_w.append(m.weight.numel())\n", + " nonzeros.append(nonzero)\n", + " k, i, f = m.get_weight_quantization_bits()\n", + " k_b, i_b, f_b = m.get_bias_quantization_bits()\n", + "\n", + " integer_bits.append(i)\n", + " integer_bias_bits.append(i)\n", + " fractional_bias_bits.append(f)\n", + "\n", + " #nonzeros = np.ndarray(nonzeros)\n", + " #total_w = np.ndarray(total_w)\n", + "\n", + " total_remaining_ratio = sum(nonzeros) / sum(total_w)\n", + "\n", + " total_n_weights = sum(p.numel() for p in model.parameters() if p.requires_grad)\n", + "\n", + " fig1, ax = plt.subplots(1, 2, figsize=(10, 6))\n", + "\n", + " ax[0].barh(range(len(names)), remaining)\n", + " ax[0].set_yticks(range(len(names)))\n", + " ax[0].set_yticklabels(names)\n", + " ax[0].invert_yaxis() \n", + "\n", + " new_xtick = []\n", + " for i in ax[0].get_xticklabels():\n", + " xtick = f\"{float(i.get_text()) * 100}%\"\n", + " new_xtick.append(xtick)\n", + " ax[0].set_xticklabels(new_xtick)\n", + " ax[0].title.set_text(f\"Remaining weights per layer (%) \\n Total ratio of remaining weights: {total_remaining_ratio*100:.2f} %\")\n", + "\n", + " ax[1].barh(range(len(nonzeros)), total_w, color=\"lightcoral\", label=\"pruned weights\")\n", + " ax[1].barh(range(len(nonzeros)), nonzeros, color=\"steelblue\", label=\"nonzero weights\")\n", + " ax[1].set_yticks(range(len(names)))\n", + " ax[1].set_yticklabels(names)\n", + " ax[1].invert_yaxis() # keep first layer at the top\n", + " ax[1].title.set_text(f\"Weights per layer \\n Total number of weights: {total_n_weights} \")\n", + " ax[1].legend()\n", + "\n", + " plt.tight_layout()\n", + " return fig1, ax\n", + "\n", + "\n", + "# this will be left unused in this notebook.\n", + "def plot_model_ebops(model):\n", + " fig2, ax2_1 = plt.subplots(figsize=(7, 5))\n", + " ax2_2 = ax2_1.twiny()\n", + "\n", + " # fill the metrics yourself!\n", + " reference_ebops = 8.9e9\n", + " pruned_ebops = get_ebops(model)\n", + " pruned_quantized_smaller_architecture_ebops = 2.02e8\n", + " compressed_smaller_architecture_ebops_2 = 3e7\n", + " fpga_ebops = 350_000 # doi.org/10.22323/1.485.0081\n", + "\n", + "\n", + " labels = [\"Uncompressed\", \"Pruned \\n (DST)\", \"Pruned \\n (DST) \\n & \\n Quantized \\n (10-bit data, \\n 6-bit weights)\"]\n", + " values = [reference_ebops, pruned_ebops, pruned_quantized_smaller_architecture_ebops]\n", + " accuracys = [71.6, 68.7, 65.2]\n", + " x = np.arange(len(labels))\n", + " width = 0.5\n", + "\n", + " ax2_2.set_xscale(\"log\")\n", + " ax2_2.set_ylabel(\"Models\")\n", + " ax2_2.set_xlabel(\"EBOPs (log.)\")\n", + "\n", + " h = 0.35\n", + " ax2_1.barh(x + h/2, accuracys, height=h, color=\"red\", label=\"Accuracy\")\n", + " ax2_2.barh(x - h/2, values, height=h, color=\"blue\", label=\"EBOPs\")\n", + "\n", + " \n", + " ax2_1.set_xlabel(\"Accuracy (%)\")\n", + " ax2_1.set_xlim(0, 100)\n", + " ax2_2.axvline(fpga_ebops, linestyle='dashed', color=\"navy\", label='EBOPs on FPGAs')\n", + "\n", + " handles1, labels1 = ax2_2.get_legend_handles_labels()\n", + " handles2, labels2 = ax2_1.get_legend_handles_labels()\n", + " ax2_2.legend(handles1 + handles2, labels1 + labels2, loc=\"best\")\n", + "\n", + " ax2_2.set_yticks(x)\n", + " ax2_2.set_yticklabels(labels, rotation=0, va=\"center\")\n", + " ax2_2.invert_yaxis() # \"Reference\" on top, matching original left-to-right order\n", + "\n", + " plt.suptitle(f\"EBOPs comparison of compressed classifier\")\n", + "\n", + " plt.tight_layout()\n", + " return fig2, (ax2_1, ax2_2)\n", + "\n" + ] + }, + { + "cell_type": "markdown", + "id": "b6efe3c9", + "metadata": {}, + "source": [ + "The original model is trained on the JetSet dataset: [mc-flavtag-ttbar-small.h5](https://opendata.cern.ch/record/93940) , in which the \"light\" class is overrepresents. Hence, it is resampled into a 100k- and 1000k-sample dataset with equal amount of samples per class, which we load into the heptoken.data.SingleFileMapModule class. If you need to resample yourself, use view_data.ipynb." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "18995b80", + "metadata": {}, + "outputs": [], + "source": [ + "data_path = Path(\"/shared/data/ttbar-1000k_balanced.h5\")" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "55467460", + "metadata": {}, + "outputs": [], + "source": [ + "# loading 1M samples\n", + "n_points = 1_000_000\n", + "dataset = SingleFileMapModule(\n", + " data_path=data_path,\n", + " batch_size=1024,\n", + " n_classes=4,\n", + " num_workers=1,\n", + " persistent_workers=True,\n", + " jet_features=JET_FEATURES,\n", + " cst_features=TRACK_FEATURES,\n", + " max_jet_pt=7_000_000,\n", + " max_cst_pt=1_000_000,\n", + " num_jets=n_points,\n", + " num_csts=n_points\n", + " )\n" + ] + }, + { + "cell_type": "markdown", + "id": "db33c7d0", + "metadata": {}, + "source": [ + "This is a LLM-inspired function that analyzes the dataset for us." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "66eddd7e", + "metadata": {}, + "outputs": [], + "source": [ + "def mapdataset_stats(split):\n", + " \"\"\"\n", + " Compute track-feature statistics and label counts for a split returned by\n", + " SingleFileMapModule: dm.train_set, dm.valid_set, or dm.test_set.\n", + " \"\"\"\n", + " ds = split.dataset # underlying MapDataset\n", + " idx = np.asarray(split.indices)\n", + "\n", + " csts = ds.data_dict[\"csts\"][idx] # [n_jets, n_csts, n_features]\n", + " mask = ds.data_dict[\"mask\"][idx].astype(bool) # [n_jets, n_csts]\n", + " labels = ds.data_dict[\"labels\"][idx]\n", + "\n", + " # Select only valid constituents, flatten to [n_valid_tracks, n_features].\n", + " valid_csts = csts[mask]\n", + "\n", + " if valid_csts.shape[0] == 0:\n", + " raise ValueError(\"This split contains no valid constituents.\")\n", + "\n", + " mean = valid_csts.mean(axis=0)\n", + " std = valid_csts.std(axis=0)\n", + "\n", + " # rowvar=False: features are variables/columns\n", + " cov = np.cov(valid_csts, rowvar=False)\n", + " corrcoef = np.corrcoef(valid_csts, rowvar=False)\n", + " return {\n", + " \"n_jets\": len(idx),\n", + " \"n_tracks\": int(mask.sum()),\n", + " \"mean\": mean,\n", + " \"std\": std,\n", + " \"cov\": cov,\n", + " \"corr\": corrcoef,\n", + " \"bincount\": np.bincount(labels, minlength=4),\n", + " \"labels\": labels,\n", + " }" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "0745ebb1", + "metadata": {}, + "outputs": [], + "source": [ + "dataset.setup(\"fit\") # harmless for SingleFileMapModule; splits already exist\n", + "\n", + "train_stats = mapdataset_stats(dataset.train_set)\n", + "val_stats = mapdataset_stats(dataset.valid_set)\n", + "test_stats = mapdataset_stats(dataset.test_set)" + ] + }, + { + "cell_type": "markdown", + "id": "72d5a6a1", + "metadata": {}, + "source": [ + "This computes the metrics and plot them in the next cell. We see the magnitude of each feature (=input dimension), the distribution of the classes and the correlation matrix." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "d708188f", + "metadata": {}, + "outputs": [], + "source": [ + "for name, s in [(\"TRAIN\", train_stats), (\"VAL\", val_stats)]:\n", + " print(f\"\\n=== {name} ===\")\n", + " print(\"n_jets:\", s[\"n_jets\"], \" n_valid_tracks:\", s[\"n_tracks\"])\n", + " for feat, m, sd in zip(TRACK_FEATURES, s[\"mean\"], s[\"std\"]):\n", + " print(f\"{feat:45s} mean={m:.4f} std={sd:.4f}\")\n", + " print(\"corr shape:\", s[\"corr\"].shape)\n", + " print(\"bincount (light, c, b, tau):\", s[\"bincount\"].tolist())" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "9492cb5d", + "metadata": {}, + "outputs": [], + "source": [ + "x = np.arange(len(TRACK_FEATURES))\n", + "width = 0.35\n", + "\n", + "fig, ax = plt.subplots(figsize=(14, 5))\n", + "ax.bar(x - width/2, train_stats[\"mean\"], width, yerr=train_stats[\"std\"], label=\"Train\", capsize=3)\n", + "ax.bar(x + width/2, val_stats[\"mean\"], width, yerr=val_stats[\"std\"], label=\"Val\", capsize=3)\n", + "ax.set_xticks(x)\n", + "ax.set_xticklabels(TRACK_FEATURES, rotation=60, ha=\"right\")\n", + "ax.set_ylabel(\"Value\")\n", + "ax.set_title(\"Feature mean/std: train vs validation\")\n", + "ax.legend()\n", + "plt.tight_layout()\n", + "\n", + "plt.show()\n", + "\n", + "fig, ax = plt.subplots(figsize=(6, 4))\n", + "xc = np.arange(len(CLASS_NAMES))\n", + "ax.bar(xc - width/2, train_stats[\"bincount\"], width, label=\"Train\")\n", + "ax.bar(xc + width/2, val_stats[\"bincount\"], width, label=\"Val\")\n", + "ax.set_yscale(\"log\")\n", + "ax.set_xticks(xc)\n", + "ax.set_xticklabels(CLASS_NAMES)\n", + "ax.set_ylabel(\"Jet count (log)\")\n", + "ax.set_title(\"Jet flavor class counts: train vs validation\")\n", + "ax.legend()\n", + "plt.tight_layout()\n", + "\n", + "plt.show()\n", + "\n", + "fig, ax = plt.subplots(figsize=(8, 7))\n", + "im = ax.imshow(train_stats[\"corr\"], cmap=\"RdBu\",\n", + " vmin=-np.abs(train_stats[\"corr\"]).max(), vmax=np.abs(train_stats[\"corr\"]).max())\n", + "ax.set_xticks(np.arange(len(TRACK_FEATURES)))\n", + "ax.set_yticks(np.arange(len(TRACK_FEATURES)))\n", + "ax.set_xticklabels(TRACK_FEATURES, rotation=60, ha=\"right\", fontsize=8)\n", + "ax.set_yticklabels(TRACK_FEATURES, fontsize=8)\n", + "ax.set_title(\"Correlation matrix (train split)\")\n", + "fig.colorbar(im, ax=ax)\n", + "plt.tight_layout()\n", + "\n", + "plt.show()" + ] + }, + { + "cell_type": "markdown", + "id": "3f528d8f", + "metadata": {}, + "source": [ + "We now construct the machine learning model from feature_clf_model, using the class FeatureClassifier. The layers in token_classifier.py and transformer.py are exchanged from torch.nn layers to PQuantML layers.\n", + "The files are pasted in the next two cells.\n", + "\n", + "This is essentially the architecture we use:\n", + "\n", + "![Architecture](/shared/figs/arch.png)" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "6025b3c5", + "metadata": {}, + "outputs": [], + "source": [ + "\"\"\"transformer.py\n", + "Transformer building block for sequence-aware encoding.\n", + "\n", + "Extracted verbatim from heptokens.models.transformer (only the ``Transformer``\n", + "class is needed by the FeatureClassifier).\n", + "\"\"\"\n", + "\n", + "def make_data_quantizer(config) -> Quantizer:\n", + " \"\"\"Build a per-tensor data-lane quantizer from a pquant config.\n", + "\n", + " Matches how pquant's automatic pass constructs data-edge quantizers\n", + " (see pquant.core.torch.tracing._insert_missing_quantizers), so manually\n", + " inserted activation quantizers are consistent with the auto-inserted ones.\n", + " \"\"\"\n", + " qp = config.quantization_parameters\n", + " return Quantizer(\n", + " k=qp.default_data_keep_negatives,\n", + " i=qp.default_data_integer_bits,\n", + " f=qp.default_data_fractional_bits,\n", + " overflow=qp.overflow_mode_data,\n", + " round_mode=qp.round_mode,\n", + " is_heterogeneous=False,\n", + " is_data=True,\n", + " granularity=\"per_tensor\",\n", + " hgq_gamma=qp.hgq_gamma,\n", + " )\n", + "\n", + "class PQTransformerEncoder(nn.Module):\n", + " \"\"\"\n", + " PQTransformerEncoder\n", + " \"\"\"\n", + "\n", + " def __init__(\n", + " self,\n", + " config,\n", + " num_layers: int,\n", + " d_model: int,\n", + " nhead: int,\n", + " dim_feedforward: int = 2048,\n", + " activation=\"relu\",\n", + " layer_norm_eps: float = 1e-5,\n", + " norm_first: bool = False,\n", + " bias: bool = True,\n", + " norm: nn.Module | None = None,\n", + " ) -> None:\n", + " super().__init__()\n", + " self.layers = nn.ModuleList(\n", + " [\n", + " PQTransformerEncoderLayer(\n", + " config,\n", + " d_model=d_model,\n", + " nhead=nhead,\n", + " dim_feedforward=dim_feedforward,\n", + " activation=activation,\n", + " layer_norm_eps=layer_norm_eps,\n", + " norm_first=norm_first,\n", + " bias=bias,\n", + " )\n", + " for _ in range(num_layers)\n", + " ]\n", + " )\n", + " self.num_layers = num_layers\n", + " self.norm = norm\n", + "\n", + " def forward(\n", + " self,\n", + " src: torch.Tensor,\n", + " mask: torch.Tensor | None = None,\n", + " src_key_padding_mask: torch.Tensor | None = None,\n", + " ) -> torch.Tensor:\n", + " output = src\n", + " for mod in self.layers:\n", + " output = mod(\n", + " output,\n", + " src_mask=mask,\n", + " src_key_padding_mask=src_key_padding_mask,\n", + " )\n", + " #print(f\"TransformerEncoder, output {output}\")\n", + " if self.norm is not None:\n", + " output = self.norm(output)\n", + " return output\n", + " \n", + "\n", + "class PQTransformerEncoderLayer(nn.Module):\n", + " r\"\"\"\n", + " PQ'd version of torch.nn.TransformerEncoderLayer\n", + " TransformerEncoderLayer is made up of self-attn and feedforward network.\n", + "\n", + " \"\"\"\n", + "\n", + " __constants__ = [\"norm_first\"]\n", + "\n", + " def __init__(\n", + " self,\n", + " config,\n", + " d_model: int,\n", + " nhead: int,\n", + " dim_feedforward: int = 2048,\n", + " activation: str | Callable[[torch.Tensor], torch.Tensor] = \"relu\",\n", + " layer_norm_eps: float = 1e-5,\n", + " batch_first: bool = False,\n", + " norm_first: bool = False,\n", + " bias: bool = True,\n", + " ) -> None:\n", + " super().__init__()\n", + " self.self_attn = PQMultiheadAttention(\n", + " config,\n", + " embed_dim=d_model,\n", + " num_heads=nhead,\n", + " bias=bias,\n", + " batch_first=True,\n", + " quantize_input=True,\n", + " quantize_output=True,\n", + " )\n", + "\n", + " self.input_quantizer = make_data_quantizer(config)\n", + "\n", + " # Feed-forward network.\n", + " self.linear1 = PQDense(config, d_model, dim_feedforward, bias=bias, quantize_output=False) # false because it goes into self.activation either way\n", + " self.linear2 = PQDense(config, dim_feedforward, d_model, bias=bias, quantize_output=True)\n", + "\n", + " self.norm_first = norm_first\n", + " self.norm1 = PQLayerNorm(config, d_model, eps=layer_norm_eps, bias=bias, quantize_output=True) # input is quantized\n", + " self.norm2 = PQLayerNorm(config, d_model, eps=layer_norm_eps, bias=bias, quantize_output=False) \n", + "\n", + " self.activation = PQActivation(config, activation, quantize_input=True) # input from linear1\n", + " \n", + "\n", + " def forward(\n", + " self,\n", + " src: torch.Tensor,\n", + " src_mask: torch.Tensor | None = None,\n", + " src_key_padding_mask: torch.Tensor | None = None,\n", + " ) -> torch.Tensor:\n", + " x = self.input_quantizer(src)\n", + " # print(f\"TransformerEncoderLayer pre {x}\")\n", + " if self.norm_first:\n", + " x = x + self._sa_block(self.norm1(x), src_mask, src_key_padding_mask)\n", + " x = x + self._ff_block(self.norm2(x))\n", + " # default\n", + " else:\n", + " x = self.norm1(x + self._sa_block(x, src_mask, src_key_padding_mask))\n", + " x = self.norm2(x + self._ff_block(x))\n", + "\n", + " # print(f\"TransformerEncoderLayer post sa/ff/norm {x}\")\n", + " return x\n", + "\n", + " # self-attention block\n", + " def _sa_block(\n", + " self,\n", + " x: torch.Tensor,\n", + " attn_mask: torch.Tensor | None,\n", + " key_padding_mask: torch.Tensor | None,\n", + " ) -> torch.Tensor:\n", + " x = self.self_attn(\n", + " x,\n", + " x,\n", + " x,\n", + " key_padding_mask=key_padding_mask,\n", + " attn_mask=attn_mask,\n", + " need_weights=False\n", + " )[0]\n", + " #print(f\"Post Self Attention {x}\")\n", + " return x\n", + "\n", + " # feed-forward block\n", + " def _ff_block(self, x: torch.Tensor) -> torch.Tensor:\n", + " #print(f\"Pre-FF (lin(act(lin(x)))) {x}\")\n", + " return self.linear2(self.activation(self.linear1(x))) \n", + " \n", + "\n", + "class Transformer(nn.Module):\n", + " \"\"\"\n", + " Transformer encoder stack with input/output projections.\n", + "\n", + " Preserves the input sequence length and supports a padding mask.\n", + " \"\"\"\n", + "\n", + " def __init__(\n", + " self,\n", + " config,\n", + " input_dim: int,\n", + " output_dim: int,\n", + " *,\n", + " d_model: int = 128,\n", + " n_heads: int = 8,\n", + " num_layers: int = 4,\n", + " dim_feedforward: int = 512,\n", + " activation = \"relu\",\n", + " ) -> None:\n", + " super().__init__()\n", + " self.input_proj = PQDense(config, input_dim, d_model)\n", + " \"\"\"\n", + " layer = PQTransformerEncoderLayer(\n", + " config,\n", + " d_model=d_model,\n", + " nhead=n_heads,\n", + " dim_feedforward=dim_feedforward,\n", + " activation=activation,\n", + " batch_first=True,\n", + " )\"\"\"\n", + " self.encoder = PQTransformerEncoder(\n", + " config, \n", + " num_layers=num_layers, \n", + " d_model=d_model, \n", + " nhead=n_heads,\n", + " dim_feedforward=dim_feedforward,\n", + " activation=activation\n", + " )\n", + " self.output_proj = PQDense(config, d_model, output_dim, quantize_output=True)\n", + "\n", + " def forward(self, x: torch.Tensor, mask: torch.Tensor | None = None) -> torch.Tensor:\n", + " \"\"\"Forward pass through the transformer.\n", + "\n", + " Args:\n", + " x: Tensor of shape [batch, n_csts, input_dim].\n", + " mask: Boolean tensor of shape [batch, n_csts] where True means valid.\n", + "\n", + " Returns:\n", + " Tensor of shape [batch, n_csts, output_dim].\n", + " \"\"\"\n", + " squeeze_batch = False\n", + " # removed for PQ\n", + " \"\"\"\n", + " if x.dim() == 2:\n", + " x = x.unsqueeze(0)\n", + " squeeze_batch = True\n", + " \"\"\"\n", + "\n", + " key_padding_mask = None\n", + " if mask is not None:\n", + " # mask = torch.tensor(mask, dtype=torch.bool)# removed for ONNX conversion, assuming mask is boolean tensor\n", + " # removed for PQ\n", + " \"\"\"\n", + " if mask.dim() == 1:\n", + " mask = mask.unsqueeze(0)\n", + " \"\"\"\n", + " # Zero out padded positions to avoid NaNs in attention/FFN paths.\n", + " \n", + " x *= mask.to(torch.float32).unsqueeze(-1)\n", + " key_padding_mask = ~mask\n", + "\n", + " # Remove for PQ\n", + " \"\"\"\n", + " if empty_sequences.any():\n", + " mask_for_encoder = mask.clone()\n", + " # Ensure at least one valid token so attention doesn't see all padding.\n", + " mask_for_encoder[empty_sequences, 0] = True\n", + " key_padding_mask = ~mask_for_encoder\n", + " x[empty_sequences] = 0.0\n", + " else:\n", + " key_padding_mask = ~mask\n", + " \"\"\"\n", + " x = self.input_proj(x)\n", + " #print(\"Transformer input proj\")\n", + " #print(x)\n", + " x = self.encoder(x, src_key_padding_mask=key_padding_mask)\n", + " #print(f\"Transformer encoder {x} with key padding mask {key_padding_mask}\")\n", + "\n", + " x = self.output_proj(x)\n", + "\n", + " if mask is not None:\n", + " x *= mask.to(torch.float32).unsqueeze(-1)\n", + " # Remove for PQ\n", + " \"\"\"\n", + " if empty_sequences is not None and empty_sequences.any():\n", + " x[empty_sequences] = 0.0\n", + " \"\"\"\n", + "\n", + " if squeeze_batch:\n", + " x = x.squeeze(0)\n", + "\n", + " #print(f\"Output projection {x}\")\n", + " return x\n" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "c92acb16", + "metadata": {}, + "outputs": [], + "source": [ + "\"\"\"token_classifier.py\n", + "Standalone FeatureClassifier (and its base JetClassifier).\n", + "\n", + "Extracted from heptokens.models.token_classifier. Only the classes reachable\n", + "from ``FeatureClassifier`` are kept:\n", + "\n", + " FeatureClassifier -> JetClassifier -> {FeatureEmbedder, TransformerEncoder,\n", + " MeanPooler / MaxPooler / ClsTokenPooler}\n", + "\n", + "The VQ-VAE-dependent embedders (TokenEmbedder, VectorEmbedder) and classifiers\n", + "(TokenClassifier, VectorClassifier) are intentionally omitted -- they require a\n", + "tokenizer checkpoint and are not part of this handoff.\n", + "\n", + "Architecture (FeatureClassifier):\n", + " {\"csts\": (B, N, 20), \"mask\": (B, N) bool}\n", + " -> FeatureEmbedder (PQDense 20 -> d_model)\n", + " -> TransformerEncoder (Transformer encoder stack, mean/max/cls)\n", + " -> Pooler -> PQDense(d_model -> n_classes) -> logits (B, n_classes)\n", + "\"\"\"\n", + "\n", + "\n", + "class TransformerEncoder(SequenceEncoder):\n", + " \"\"\"Transformer-based sequence encoder.\"\"\"\n", + "\n", + " def __init__(\n", + " self,\n", + " config,\n", + " input_dim: int,\n", + " d_model: int,\n", + " n_heads: int = 8,\n", + " num_layers: int = 4,\n", + " dim_feedforward: int = 512,\n", + " activation = \"relu\",\n", + " ) -> None:\n", + " super().__init__()\n", + " self.d_model = d_model\n", + " self.transformer = Transformer(\n", + " config,\n", + " input_dim=d_model,\n", + " output_dim=d_model,\n", + " d_model=d_model,\n", + " n_heads=n_heads,\n", + " num_layers=num_layers,\n", + " dim_feedforward=dim_feedforward,\n", + " activation=activation,\n", + " )\n", + "\n", + " @property\n", + " def output_dim(self) -> int:\n", + " return self.d_model\n", + "\n", + " def encode(self, x: T.Tensor, mask: T.BoolTensor) -> T.Tensor:\n", + " return self.transformer(x, mask=mask)\n", + "\n", + "\n", + "class MeanPooler(Pooler):\n", + " \"\"\"Mean pooling over valid positions.\"\"\"\n", + "\n", + " def __init__(self, config) -> None:\n", + " super().__init__()\n", + " # Quantize the genuine data operands of the masked mean: the input\n", + " # activation, the accumulated sum, and the per-jet count (the divisor).\n", + " # The mask itself is a structural {0, 1} selector, not a data activation,\n", + " # so it is left unquantized.\n", + " self.x_quantizer = make_data_quantizer(config)\n", + " self.sum_quantizer = make_data_quantizer(config)\n", + " self.count_quantizer = make_data_quantizer(config)\n", + " self.num_csts = config.training_parameters.num_csts\n", + "\n", + " def pool(self, x: T.Tensor, mask: T.BoolTensor) -> T.Tensor:\n", + " x = self.x_quantizer(x) \n", + " mask_f = mask.to(T.float32)\n", + " valid_sum = self.sum_quantizer((x * mask_f.unsqueeze(-1)).sum(dim=1)) \n", + " valid_count = mask_f.sum(dim=1, keepdim=True).clamp(min=1) # no quantizer because mask is 0/1\n", + " \n", + " valid_ratio = self.count_quantizer(valid_sum / valid_count) # quantizer because this is a ratio\n", + " return valid_ratio\n", + " \n", + "class MaxPooler(Pooler):\n", + " \"\"\"Max pooling over valid positions.\"\"\"\n", + "\n", + " def pool(self, x: T.Tensor, mask: T.BoolTensor) -> T.Tensor:\n", + " x_masked = x.masked_fill(~mask.unsqueeze(-1), float(\"-inf\"))\n", + " return x_masked.max(dim=1)[0]\n", + "\n", + "class FeatureEmbedder(Embedder):\n", + " \"\"\"Embedder for raw constituent features.\"\"\"\n", + "\n", + " def __init__(self, config, input_dim: int, d_model: int) -> None:\n", + " super().__init__()\n", + " self.input_dim = input_dim\n", + " self.d_model = d_model\n", + " self.projection = PQDense(config, input_dim, d_model, quantize_input=True)\n", + "\n", + " @property\n", + " def output_dim(self) -> int:\n", + " return self.d_model\n", + "\n", + " def embed(self, csts, mask) -> T.Tensor:\n", + " csts = T.nan_to_num(csts, nan=0.0)\n", + " # Zero out masked positions\n", + " mask_f = mask.to(T.float32)\n", + " mask_expanded = mask_f.unsqueeze(-1)\n", + " csts = csts * mask_expanded\n", + " \n", + " #csts = T.where(mask.unsqueeze(-1), csts, T.zeros_like(csts))\n", + "\n", + " # embeddings = self.projection(csts) put directly in mask argument\n", + " return self.projection(csts)\n", + "\n", + "\n", + "class JetClassifier(ScheduledOptimiserMixin, LightningModule):\n", + " \"\"\"General-purpose jet classifier with pluggable components.\n", + "\n", + " Architecture:\n", + " Input -> Embedder -> Encoder -> Pooler -> Head -> Logits\n", + " \"\"\"\n", + "\n", + " def __init__(\n", + " self,\n", + " *,\n", + " config,\n", + " embedder: Embedder,\n", + " encoder: SequenceEncoder,\n", + " pooler: Pooler,\n", + " n_classes: int,\n", + " learning_rate: float = 1e-3,\n", + " **kwargs,\n", + " ) -> None:\n", + " super().__init__()\n", + " self.save_hyperparameters(ignore=[\"embedder\", \"encoder\", \"pooler\", \"layers\"])\n", + "\n", + " self.F = config.training_parameters.num_features # 20\n", + " self.N = config.training_parameters.num_csts # 40\n", + "\n", + " self.input_shape = (self.N, self.F+1)\n", + "\n", + " # remove this because this is only valid for original implementation, but we replace it with a PQ'd version\n", + " \"\"\"# Validate compatibility\n", + " if embedder.output_dim != encoder.encode.__code__.co_varnames[1:2][0]: # Quick check\n", + " log.warning(\n", + " f\"Embedder output_dim ({embedder.output_dim}) may not match \"\n", + " f\"encoder input_dim. Check compatibility.\"\n", + " )\"\"\"\n", + "\n", + " self.embedder = embedder\n", + " self.encoder = encoder\n", + " self.pooler = pooler\n", + " self.n_classes = n_classes\n", + "\n", + " self.config = config\n", + " self.training_config = self.config.training_parameters\n", + "\n", + " # Classification head\n", + " self.classifier = PQDense(config, encoder.output_dim, n_classes, quantize_output=True)\n", + "\n", + " # Metrics\n", + " self.train_acc = Accuracy(\"multiclass\", num_classes=n_classes, average=\"macro\")\n", + " self.valid_acc = Accuracy(\"multiclass\", num_classes=n_classes, average=\"macro\")\n", + " self.test_acc = Accuracy(\"multiclass\", num_classes=n_classes, average=\"macro\")\n", + "\n", + " self.train_acc_classwise = Accuracy(\"multiclass\", num_classes=n_classes, average=\"none\")\n", + " self.valid_acc_classwise = Accuracy(\"multiclass\", num_classes=n_classes, average=\"none\")\n", + " self.test_acc_classwise = Accuracy(\"multiclass\", num_classes=n_classes, average=\"none\")\n", + "\n", + " # AUC metrics (one-vs-rest for multiclass)\n", + " self.train_auc = AUROC(task=\"multiclass\", num_classes=n_classes, average=\"macro\")\n", + " self.valid_auc = AUROC(task=\"multiclass\", num_classes=n_classes, average=\"macro\")\n", + " self.test_auc = AUROC(task=\"multiclass\", num_classes=n_classes, average=\"macro\")\n", + "\n", + " \n", + " # F1 metrics (one-vs-rest for multiclass)\n", + " self.train_f1 = F1Score(task=\"multiclass\", num_classes=n_classes, average=\"macro\")\n", + " self.valid_f1 = F1Score(task=\"multiclass\", num_classes=n_classes, average=\"macro\")\n", + " self.test_f1 = F1Score(task=\"multiclass\", num_classes=n_classes, average=\"macro\")\n", + "\n", + " self.valid_confusion_matrix = ConfusionMatrix(task=\"multiclass\", num_classes=n_classes, normalize=\"true\")\n", + " self.test_confusion_matrix = ConfusionMatrix(task=\"multiclass\", num_classes=n_classes, normalize=\"true\")\n", + "\n", + " self.class_weights = T.ones(n_classes)\n", + "\n", + " # Store outputs for ROC plotting\n", + " self.validation_outputs = []\n", + "\n", + " def forward(self, input: T.Tensor) -> T.Tensor:\n", + " \"\"\"Single-tensor version instead of dict version\"\"\"\n", + " csts = input[:, :, :self.F]\n", + " mask = input[:, :, self.F] > 0 # implicit conversion back from csts.dtype to bool\n", + "\n", + " x = self.embedder.embed(csts, mask)\n", + " x = self.encoder.encode(x, mask)\n", + " x = self.pooler.pool(x, mask)\n", + " return self.classifier(x)\n", + " \n", + " def pq_loss(self, output, labels, loss_fn=cross_entropy):\n", + " loss = loss_fn(output, labels, label_smoothing=0.01).to(output.device)\n", + " compression_loss = get_model_losses(self, T.tensor(0.).to(output.device))\n", + " return loss + compression_loss\n", + "\n", + " def _shared_step(self, batch: dict, prefix: str) -> T.Tensor:\n", + " x, labels = batch\n", + " x = x.to(device=self.device)\n", + " labels = labels.to(device=self.device)\n", + "\n", + " output = self.forward(x)\n", + "\n", + " loss = self.pq_loss(output, labels) \n", + "\n", + " self.log(f\"{prefix}/total_loss\", loss)\n", + "\n", + " acc = getattr(self, f\"{prefix}_acc\")\n", + " acc(output, labels)\n", + " self.log(f\"{prefix}/acc\", acc)\n", + "\n", + " acc_classwise = getattr(self, f\"{prefix}_acc_classwise\")\n", + " acc_per_class = acc_classwise(output, labels)\n", + " for i, acc in enumerate(acc_per_class):\n", + " self.log(f\"{prefix}/acc_class_{i}_{CLASS_NAMES[i]}\", acc, on_epoch=True)\n", + "\n", + " auc = getattr(self, f\"{prefix}_auc\")\n", + " probs = T.softmax(output, dim=1)\n", + " auc(probs, labels)\n", + " self.log(f\"{prefix}/auc\", auc)\n", + "\n", + " f1 = getattr(self, f\"{prefix}_f1\")\n", + " f1(probs, labels)\n", + " self.log(f\"{prefix}/f1\", f1)\n", + "\n", + " if prefix == \"test\":\n", + " self.test_confusion_matrix(probs, labels)\n", + " elif prefix == \"valid\":\n", + " self.valid_confusion_matrix(probs, labels)\n", + " return loss\n", + "\n", + " def training_step(self, batch_dict: dict) -> T.Tensor:\n", + " return self._shared_step(batch_dict, \"train\")\n", + "\n", + " def validation_step(self, batch_dict: dict) -> T.Tensor:\n", + " return self._shared_step(batch_dict, \"valid\")\n", + "\n", + " def predict_step(self, batch: dict) -> dict:\n", + " x, labels = batch\n", + "\n", + " output = self.forward(x)\n", + "\n", + " return {\"output\": output, \"label\": labels.unsqueeze(-1)}\n", + " \n", + " def test_step(self, batch: dict) -> dict:\n", + " return self._shared_step(batch, \"test\")\n", + " \n", + " def on_train_epoch_start(self):\n", + " if self.current_epoch == 0:\n", + " pass # initial\n", + " self.train()\n", + " pre_epoch_functions(self, self.current_epoch, self.training_config.pretraining_epochs)\n", + " if self.training_config.rounds == 0:\n", + " if (self.current_epoch == self.training_config.pretraining_epochs + self.training_config.epochs - 1) and (self.training_config.fine_tuning_epochs != 0):\n", + " print(\"Finetuning starting\")\n", + " pre_finetune_functions(self)\n", + " else:\n", + " if (self.current_epoch == self.training_config.pretraining_epochs + self.training_config.rounds * self.training_config.epochs - 1) and (self.training_config.fine_tuning_epochs != 0):\n", + " print(\"Finetuning starting\")\n", + " pre_finetune_functions(self)\n", + " \n", + " def on_train_epoch_end(self):\n", + " # general post epoch function\n", + " post_epoch_functions(self, self.current_epoch, self.training_config.pretraining_epochs)\n", + " \n", + " def on_validation_epoch_start(self):\n", + " self.eval()\n", + " \n", + " def on_validation_epoch_end(self):\n", + " self.to(self.device)\n", + " self.log(f\"valid/remaining_weights\", get_layer_keep_ratio(self))\n", + " self.log(f\"valid/EBOPs\", get_ebops(self))\n", + "\n", + " conf_matrix_plot, conf_matrix_ax = self.valid_confusion_matrix.plot(\n", + " labels=CLASS_NAMES\n", + " )\n", + " conf_matrix_plot.set_dpi(600)\n", + "\n", + " if self.logger is not None and hasattr(self.logger.experiment, \"add_figure\"):\n", + " self.logger.experiment.add_figure(\"valid/confusion_matrix\", conf_matrix_plot, self.current_epoch)\n", + " else:\n", + " print(\"No logger configured. Image is deleted\")\n", + " \n", + " plt.close(conf_matrix_plot)\n", + "\n", + " self.valid_confusion_matrix.reset()\n", + "\n", + " # post pretrain\n", + " if (self.current_epoch == self.training_config.pretraining_epochs - 1) or (self.training_config.pretraining_epochs == 0 and self.current_epoch == 0):\n", + " try:\n", + " print(f\"{self.trainer.train_dataloader=}\")\n", + " except Exception as e:\n", + " print(e)\n", + " train_dl = self.trainer.train_dataloader if hasattr(self, 'trainer') else None\n", + " with T.enable_grad():\n", + " post_pretrain_functions(\n", + " self, self.config, input_shape=self.input_shape, train_loader=train_dl, loss_function=self.pq_loss\n", + " ) # after validation\n", + " print(\"Pretraining ended\")\n", + "\n", + " def on_test_epoch_end(self):\n", + " conf_matrix_plot, conf_matrix_ax = self.test_confusion_matrix.plot(\n", + " labels=CLASS_NAMES\n", + " )\n", + " conf_matrix_plot.tight_layout()\n", + "\n", + " if self.logger is not None and hasattr(self.logger.experiment, \"add_figure\"):\n", + " self.logger.experiment.add_figure(\"valid/confusion_matrix\", conf_matrix_plot, self.current_epoch)\n", + " else:\n", + " print(\"No logger configured. Image is deleted\")\n", + " \n", + " plt.close(conf_matrix_plot)\n", + "\n", + " self.test_confusion_matrix.reset()\n", + "\n", + "\n", + "class FeatureClassifier(JetClassifier):\n", + " \"\"\"Explicit instantiation of JetClassifier using raw feature embedder.\"\"\"\n", + "\n", + " def __init__(\n", + " self,\n", + " *,\n", + " config,\n", + " data_sample: dict | None = None,\n", + " n_classes: int,\n", + " d_model: int = 128,\n", + " n_heads: int = 8,\n", + " num_layers: int = 4,\n", + " dim_feedforward: int = 512,\n", + " activation: str = \"relu\",\n", + " pooling: Literal[\"mean\", \"max\", \"cls\"] = \"mean\",\n", + " learning_rate: float = 1e-3,\n", + " input_dim: int | None = None,\n", + " **kwargs,\n", + " ) -> None:\n", + " # The original model reads the input feature dim from a data sample.\n", + " # We also accept an explicit ``input_dim`` so the model can be built\n", + " # without a data sample (the checkpoint does store ``data_sample``).\n", + " if data_sample is not None:\n", + " input_dim = data_sample[\"csts\"].shape[-1]\n", + " if input_dim is None:\n", + " raise ValueError(\n", + " \"Provide either `data_sample` (dict with 'csts') or `input_dim`.\"\n", + " )\n", + "\n", + " embedder = FeatureEmbedder(\n", + " config,\n", + " input_dim=input_dim,\n", + " d_model=d_model,\n", + " )\n", + " encoder = TransformerEncoder(\n", + " config,\n", + " input_dim,\n", + " d_model=d_model,\n", + " n_heads=n_heads,\n", + " num_layers=num_layers,\n", + " dim_feedforward=dim_feedforward,\n", + " activation=activation,\n", + " )\n", + " pooler_cls = {\n", + " \"mean\": MeanPooler,\n", + " \"max\": MaxPooler,\n", + " }[pooling]\n", + " pooler = pooler_cls(d_model, config) if pooling == \"cls\" else pooler_cls(config)\n", + "\n", + " super().__init__(\n", + " config=config,\n", + " embedder=embedder,\n", + " encoder=encoder,\n", + " pooler=pooler,\n", + " n_classes=n_classes,\n", + " learning_rate=learning_rate,\n", + " **kwargs,\n", + " )\n", + "\n", + "def warmup_cosine_scheduler(\n", + " optimizer,\n", + " warmup_epochs: int = 5,\n", + " min_lr: float = 1e-6,\n", + " model=None,\n", + " max_epochs: int = -1,\n", + "):\n", + " \"\"\"Linear warmup followed by cosine annealing, configured in epochs.\n", + "\n", + " If ``max_epochs`` is -1 (default), it is read from ``model.trainer``.\n", + " \"\"\"\n", + " if max_epochs < 1 and model is not None:\n", + " max_epochs = model.trainer.max_epochs\n", + " if max_epochs < 1:\n", + " raise ValueError(\"max_epochs must be positive (set it or pass a model with a trainer).\")\n", + "\n", + " warmup = min(warmup_epochs, max_epochs)\n", + " warmup_sched = T.optim.lr_scheduler.LinearLR(\n", + " optimizer, start_factor=1e-2, total_iters=warmup,\n", + " )\n", + " cosine_sched = T.optim.lr_scheduler.CosineAnnealingLR(\n", + " optimizer, T_max=max(max_epochs - warmup, 1), eta_min=min_lr,\n", + " )\n", + " return T.optim.lr_scheduler.SequentialLR(\n", + " optimizer,\n", + " schedulers=[warmup_sched, cosine_sched],\n", + " milestones=[warmup],\n", + " )\n", + "\n" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "e254b8be", + "metadata": {}, + "outputs": [], + "source": [ + "time_str = timestamp()\n", + "save_path = f\"/shared/logs/{time_str}\"\n", + "onnx_path = f\"{save_path}/model_last.onnx\"\n", + "\n", + "os.makedirs(save_path, exist_ok=True)\n", + "\n", + "logging.basicConfig(\n", + "level=logging.INFO,\n", + "format=\"%(asctime)s [%(levelname)s] %(message)s\",\n", + "handlers=[logging.FileHandler(f\"{save_path}/output.log\", mode='w+'), logging.StreamHandler()],\n", + ")\n" + ] + }, + { + "cell_type": "markdown", + "id": "d13ede96", + "metadata": {}, + "source": [ + "This loads preprocessing scripts from heptokens, namely normalizers" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "1d483d57", + "metadata": {}, + "outputs": [], + "source": [ + "if PREPROCESSING_ON:\n", + " jet_q_path = Path(\"/shared/pqjetset_standalone/checkpoint/jet_quantiles.joblib\")\n", + " cst_q_path = Path(\"/shared/pqjetset_standalone/checkpoint/cst_quantiles.joblib\")\n", + " jet_q = joblib.load(jet_q_path)\n", + " cst_q = joblib.load(cst_q_path)\n", + " preprocess = partial(preprocess_batch, cst_fn=cst_q, jet_fn=jet_q)\n", + "else:\n", + " preprocess = lambda x: x # do nothing" + ] + }, + { + "cell_type": "markdown", + "id": "13241e97", + "metadata": {}, + "source": [ + "We load a standard config from PQuantML for the respective pruning algorithm. We start with PDP, and adapt the preloaded values wherever required. Test out different sparsity values and unstructured/structured flags! You will notice that structured pruning only allows a fraction of the sparsity possibl with unstructured pruning.\n", + "\n", + "The training should take approx. 4 hours, but it is possible to just decrease the amount of epochs to something like 10/10/5 instead of 15/30/10." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "de71918d", + "metadata": {}, + "outputs": [], + "source": [ + "config = pdp_config() \n", + "\n", + "config.training_parameters.pretraining_epochs = 15 # training before start of pruning\n", + "config.training_parameters.epochs = 30 # training with pruning\n", + "config.training_parameters.fine_tuning_epochs = 10 # training after end of pruning, pruning mask is fixed here\n", + "max_epochs = config.training_parameters.pretraining_epochs + config.training_parameters.epochs + config.training_parameters.fine_tuning_epochs\n", + "# added manually to config for tuple input of model\n", + "config.training_parameters.num_csts = 40 # amount of constituents\n", + "config.training_parameters.num_features = 20 # amount of features\n", + "\n", + "config.quantization_parameters.enable_quantization = True\n", + "config.quantization_parameters.granularity = \"per_weight\"\n", + "# granularity -> per tensor/channel/weight \n", + "config.quantization_parameters.use_relu_multiplier = False\n", + "config.quantization_parameters.use_high_granularity_quantization = False\n", + "config.quantization_parameters.hgq_beta = 1e-13\n", + "config.quantization_parameters.hgq_gamma = 1e-6\n", + "config.quantization_parameters.overflow_mode_parameters = \"SAT\"\n", + "config.quantization_parameters.overflow_mode_data = \"SAT\"\n", + "\n", + "config.quantization_parameters.default_data_keep_negatives = 0.\n", + "config.quantization_parameters.dynamic_data_quantization = True\n", + "config.quantization_parameters.default_data_integer_bits = 0.\n", + "config.quantization_parameters.default_data_fractional_bits = 8. # give all weights to fractional by default if dynamical data quantization\n", + "\n", + "config.quantization_parameters.default_weight_keep_negatives = 1.\n", + "config.quantization_parameters.default_weight_integer_bits = 0.\n", + "config.quantization_parameters.default_weight_fractional_bits = 6. # give all weights to fractional by default if dynamical data quantization\n", + "\n", + "config.pruning_parameters.enable_pruning = True\n", + "#config.pruning_parameters.max_pruning_pct = 1.00 # DST setting\n", + "#config.pruning_parameters.alpha = 1e-3 # DST setting\n", + "config.pruning_parameters.epsilon = 0.037\n", + "config.pruning_parameters.sparsity = 0.90\n", + "config.pruning_parameters.structured_pruning = False\n", + "\n", + "config.training_parameters.num_csts = 40 \n", + "config.training_parameters.num_features = 20 " + ] + }, + { + "cell_type": "markdown", + "id": "e6726465", + "metadata": {}, + "source": [ + "The original model has a (d_model, dim_feedforward, num_layers) of (128, 512, 4), which is strongly overdimensioned. \n", + "Reducing it to (128, 128, 4) is possible without hurting performance, but accelerates training.\n", + "Reducing it until (64, 64, 3) is possible with some performance loss (f1 = 0.6, f1_original = 0.66)\n", + "Model for Sanjiban is trained with (1024, 1024, 4).\n", + "All information should also be stored in the .yaml file created by Lightning." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "3c7c2623", + "metadata": {}, + "outputs": [], + "source": [ + "n_classes = 4\n", + "d_model = 128\n", + "n_heads = 8 # does not change EBOPs!\n", + "num_layers = 4\n", + "dim_feedforward = 512\n", + "activation = 'relu' # originally gelu, but for compression, usually RELU is used because its easier\n", + "learning_rate = 1e-3\n", + "pooling = 'mean'\n", + "input_dim = config.training_parameters.num_features \n", + "\n", + "model = FeatureClassifier(\n", + " config=config, \n", + " n_classes=n_classes, \n", + " d_model=d_model,\n", + " n_heads=n_heads,\n", + " num_layers=num_layers,\n", + " dim_feedforward=dim_feedforward,\n", + " activation=activation,\n", + " learning_rate=learning_rate,\n", + " pooling=pooling,\n", + " # these are tunable hyperparameters as well, but they are also pruning/quantization-dependent and do not have that big influence compared to learning rate etc.\n", + " scheduler=partial(warmup_cosine_scheduler, warmup_epochs=3, min_lr=1e-5),\n", + " input_dim=input_dim\n", + " )\n", + "add_compression_layers(model, config)" + ] + }, + { + "cell_type": "markdown", + "id": "c758b3d6", + "metadata": {}, + "source": [ + "This searches missing quantizers and adds them automatically. \n", + "Theoretically, one could also define a pure torch/keras model and run this function to add the compressors." + ] + }, + { + "cell_type": "markdown", + "id": "fbe40a76", + "metadata": {}, + "source": [ + "We move the model entirely to device (=\"cuda\").\n", + "In the hpo loop, this needs to be called in each iteration, as otherwise, the new instance of the model would be created on cpu. Especially quantizers are prone to be reinitialized on cpu despite having been on cuda previously." + ] + }, + { + "cell_type": "markdown", + "id": "0c047332", + "metadata": {}, + "source": [ + "We make an initializing forward pass with a random input. " + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "2c605e4d", + "metadata": {}, + "outputs": [], + "source": [ + "# Pseudo forward pass initialization of network\n", + "model.eval()\n", + "B, N, F = 3, config.training_parameters.num_csts, config.training_parameters.num_features\n", + "csts = torch.randn(B, N, F).to(dtype=dtype_real)\n", + "mask = torch.ones(B, N, dtype=torch.bool)\n", + "x = torch.cat([csts, mask.unsqueeze(-1).to(csts.dtype)], dim=-1)\n", + "with torch.no_grad():\n", + " model(x)" + ] + }, + { + "cell_type": "markdown", + "id": "9567934f", + "metadata": {}, + "source": [ + "This is the model architecture. We should see the model FeatureClassifier, the submodules and for each layer, the quantizers and the pruning mask." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "bf434aa9", + "metadata": {}, + "outputs": [], + "source": [ + "print(model)" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "99144f8a", + "metadata": {}, + "outputs": [], + "source": [ + "trainer = L.Trainer(\n", + " default_root_dir=save_path,\n", + " max_epochs=max_epochs,\n", + " accelerator=\"auto\",\n", + " devices=\"auto\",\n", + " precision=\"bf16-mixed\",\n", + " val_check_interval=1.0,\n", + " callbacks=[\n", + " ModelCheckpoint(dirpath=save_path, filename=\"best\",\n", + " monitor=\"valid/total_loss\", mode=\"min\", save_top_k=1),\n", + " ModelCheckpoint(dirpath=save_path, filename=\"last\"),\n", + " LearningRateMonitor(logging_interval=\"step\"),\n", + " ]\n", + " \n", + ")" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "72af1717", + "metadata": {}, + "outputs": [], + "source": [ + "model.train()\n", + "trainer.fit(model, datamodule=dataset)" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "958536fc", + "metadata": {}, + "outputs": [], + "source": [ + "model.eval()\n", + "fig, ax = plot_remaining_weights(model)\n", + "plt.savefig(f\"{save_path}/remaining_weights_pdp.png\")\n", + "torch.save(model.state_dict(), f\"{save_path}/state_manual_pdp.pth\")\n", + "state_dict_export(model, f\"{save_path}/model_state_dict_pdp.log\") # this creates a huge text file where the state dict is logged for debugging purposes.\n", + "\n" + ] + }, + { + "cell_type": "markdown", + "id": "4c72931f", + "metadata": {}, + "source": [ + "Now we use HGQ instead of \"normal\" quantization. HGQ (High-Granularity Quantization) uses differentiable quantization to find optimal number of bits." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "207ee7d5", + "metadata": {}, + "outputs": [], + "source": [ + "config = pdp_config() \n", + "\n", + "config.training_parameters.pretraining_epochs = 15 # training before start of pruning\n", + "config.training_parameters.epochs = 30 # training with pruning\n", + "config.training_parameters.fine_tuning_epochs = 10 # training after end of pruning, pruning mask is fixed here\n", + "max_epochs = config.training_parameters.pretraining_epochs + config.training_parameters.epochs + config.training_parameters.fine_tuning_epochs\n", + "# added manually to config for tuple input of model\n", + "config.training_parameters.num_csts = 40 # amount of constituents\n", + "config.training_parameters.num_features = 20 # amount of features\n", + "\n", + "config.quantization_parameters.enable_quantization = True\n", + "config.quantization_parameters.granularity = \"per_weight\"\n", + "# granularity -> per tensor/channel/weight \n", + "config.quantization_parameters.use_relu_multiplier = False\n", + "config.quantization_parameters.use_high_granularity_quantization = True\n", + "config.quantization_parameters.hgq_beta = 1e-12\n", + "config.quantization_parameters.hgq_gamma = 1e-6\n", + "config.quantization_parameters.overflow_mode_parameters = \"SAT\"\n", + "config.quantization_parameters.overflow_mode_data = \"SAT\"\n", + "\n", + "config.quantization_parameters.default_data_keep_negatives = 0.\n", + "config.quantization_parameters.dynamic_data_quantization = True\n", + "config.quantization_parameters.default_data_integer_bits = 0.\n", + "config.quantization_parameters.default_data_fractional_bits = 8. # give all weights to fractional by default if dynamical data quantization\n", + "\n", + "config.quantization_parameters.default_weight_keep_negatives = 1.\n", + "config.quantization_parameters.default_weight_integer_bits = 0.\n", + "config.quantization_parameters.default_weight_fractional_bits = 6. # give all weights to fractional by default if dynamical data quantization\n", + "\n", + "config.pruning_parameters.enable_pruning = True\n", + "#config.pruning_parameters.max_pruning_pct = 1.00 # DST setting\n", + "#config.pruning_parameters.alpha = 1e-3 # DST setting\n", + "config.pruning_parameters.epsilon = 0.037\n", + "config.pruning_parameters.sparsity = 0.90\n", + "config.pruning_parameters.structured_pruning = False\n", + "\n", + "config.training_parameters.num_csts = 40 \n", + "config.training_parameters.num_features = 20 " + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "5edb1597", + "metadata": {}, + "outputs": [], + "source": [ + "n_classes = 4\n", + "d_model = 128\n", + "n_heads = 8 # does not change EBOPs!\n", + "num_layers = 4\n", + "dim_feedforward = 512\n", + "activation = 'relu' # originally gelu, but for compression, usually RELU is used because its easier\n", + "learning_rate = 1e-3\n", + "pooling = 'mean'\n", + "input_dim = config.training_parameters.num_features \n", + "\n", + "model = FeatureClassifier(\n", + " config=config, \n", + " n_classes=n_classes, \n", + " d_model=d_model,\n", + " n_heads=n_heads,\n", + " num_layers=num_layers,\n", + " dim_feedforward=dim_feedforward,\n", + " activation=activation,\n", + " learning_rate=learning_rate,\n", + " pooling=pooling,\n", + " # these are tunable hyperparameters as well, but they are also pruning/quantization-dependent and do not have that big influence compared to learning rate etc.\n", + " scheduler=partial(warmup_cosine_scheduler, warmup_epochs=3, min_lr=1e-5),\n", + " input_dim=input_dim\n", + " )\n", + "add_compression_layers(model, config)" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "bca5194e", + "metadata": {}, + "outputs": [], + "source": [ + "# Pseudo forward pass initialization of network\n", + "model.eval()\n", + "B, N, F = 3, config.training_parameters.num_csts, config.training_parameters.num_features\n", + "csts = torch.randn(B, N, F).to(dtype=dtype_real)\n", + "mask = torch.ones(B, N, dtype=torch.bool)\n", + "x = torch.cat([csts, mask.unsqueeze(-1).to(csts.dtype)], dim=-1)\n", + "with torch.no_grad():\n", + " model(x)" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "9227231d", + "metadata": {}, + "outputs": [], + "source": [ + "model.train()\n", + "trainer.fit(model, datamodule=dataset)" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "8a2dc3ee", + "metadata": {}, + "outputs": [], + "source": [ + "model.eval()\n", + "fig, ax = plot_remaining_weights(model)\n", + "plt.savefig(f\"{save_path}/remaining_weights_pdp_hgq.png\")\n", + "torch.save(model.state_dict(), f\"{save_path}/state_manual_pdp_hgq.pth\")\n", + "state_dict_export(model, f\"{save_path}/model_state_dict_pdp_hgq.log\") # this creates a huge text file where the state dict is logged for debugging purposes." + ] + }, + { + "cell_type": "markdown", + "id": "d9ffc90b", + "metadata": {}, + "source": [ + "We repeat this with DST pruning, where we do not set an explicit sparsity, only the \"strength\" of pruning (ratio of pruning loss vs. cross-entropy loss) and the maximum pruning ratio (to prevent the whole layer getting pruned, which is okay in case of skip connections but would possibly lead to dead weights in other layeres). We import 'dst_config' instead of 'pdp_config'.\n", + "\n", + "\\alpha = 1e-3 is very aggressive, but this is okay in our network as it is overdimensioned either way." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "744eb6b7", + "metadata": {}, + "outputs": [], + "source": [ + "config = dst_config() \n", + "\n", + "config.training_parameters.pretraining_epochs = 0 # training before start of pruning\n", + "config.training_parameters.epochs = 30 # training with pruning\n", + "config.training_parameters.fine_tuning_epochs = 0 # training after end of pruning, pruning mask is fixed here\n", + "max_epochs = config.training_parameters.pretraining_epochs + config.training_parameters.epochs + config.training_parameters.fine_tuning_epochs\n", + "# added manually to config for tuple input of model\n", + "config.training_parameters.num_csts = 40 # amount of constituents\n", + "config.training_parameters.num_features = 20 # amount of features\n", + "\n", + "config.quantization_parameters.enable_quantization = True\n", + "config.quantization_parameters.granularity = \"per_weight\"\n", + "# granularity -> per tensor/channel/weight \n", + "config.quantization_parameters.use_relu_multiplier = False\n", + "config.quantization_parameters.use_high_granularity_quantization = False\n", + "config.quantization_parameters.hgq_beta = 1e-13\n", + "config.quantization_parameters.hgq_gamma = 1e-6\n", + "config.quantization_parameters.overflow_mode_parameters = \"SAT\"\n", + "config.quantization_parameters.overflow_mode_data = \"SAT\"\n", + "\n", + "config.quantization_parameters.default_data_keep_negatives = 0.\n", + "config.quantization_parameters.dynamic_data_quantization = True\n", + "config.quantization_parameters.default_data_integer_bits = 0.\n", + "config.quantization_parameters.default_data_fractional_bits = 8. # give all weights to fractional by default if dynamical data quantization\n", + "\n", + "config.quantization_parameters.default_weight_keep_negatives = 1.\n", + "config.quantization_parameters.default_weight_integer_bits = 0.\n", + "config.quantization_parameters.default_weight_fractional_bits = 6. # give all weights to fractional by default if dynamical data quantization\n", + "\n", + "config.pruning_parameters.enable_pruning = True\n", + "config.pruning_parameters.max_pruning_pct = 1.00 # DST setting\n", + "config.pruning_parameters.alpha = 1e-3 # DST setting\n", + "#config.pruning_parameters.epsilon = 0.037\n", + "#config.pruning_parameters.sparsity = 0.90\n", + "#config.pruning_parameters.structured_pruning = False\n", + "\n", + "config.training_parameters.num_csts = 40 \n", + "config.training_parameters.num_features = 20 " + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "8ab84e87", + "metadata": {}, + "outputs": [], + "source": [ + "n_classes = 4\n", + "d_model = 128\n", + "n_heads = 8 # does not change EBOPs!\n", + "num_layers = 4\n", + "dim_feedforward = 512\n", + "activation = 'relu' # originally gelu, but for compression, usually RELU is used because its easier\n", + "learning_rate = 1e-3\n", + "pooling = 'mean'\n", + "input_dim = config.training_parameters.num_features \n", + "\n", + "model = FeatureClassifier(\n", + " config=config, \n", + " n_classes=n_classes, \n", + " d_model=d_model,\n", + " n_heads=n_heads,\n", + " num_layers=num_layers,\n", + " dim_feedforward=dim_feedforward,\n", + " activation=activation,\n", + " learning_rate=learning_rate,\n", + " pooling=pooling,\n", + " # these are tunable hyperparameters as well, but they are also pruning/quantization-dependent and do not have that big influence compared to learning rate etc.\n", + " scheduler=partial(warmup_cosine_scheduler, warmup_epochs=3, min_lr=1e-5),\n", + " input_dim=input_dim\n", + " )\n", + "add_compression_layers(model, config)" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "9cabc758", + "metadata": {}, + "outputs": [], + "source": [ + "# Pseudo forward pass initialization of network\n", + "model.eval()\n", + "B, N, F = 3, config.training_parameters.num_csts, config.training_parameters.num_features\n", + "csts = torch.randn(B, N, F).to(dtype=dtype_real)\n", + "mask = torch.ones(B, N, dtype=torch.bool)\n", + "x = torch.cat([csts, mask.unsqueeze(-1).to(csts.dtype)], dim=-1)\n", + "with torch.no_grad():\n", + " model(x)" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "8db286f2", + "metadata": {}, + "outputs": [], + "source": [ + "model.train()\n", + "trainer.fit(model, datamodule=dataset)" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "d7f95c27", + "metadata": {}, + "outputs": [], + "source": [ + "model.eval()\n", + "fig, ax = plot_remaining_weights(model)\n", + "plt.savefig(f\"{save_path}/remaining_weights_dst.png\")\n", + "torch.save(model.state_dict(), f\"{save_path}/state_manual_dst.pth\")\n", + "state_dict_export(model, f\"{save_path}/model_state_dict_dst.log\") # this creates a huge text file where the state dict is logged for debugging purposes." + ] + }, + { + "cell_type": "markdown", + "id": "79446da6", + "metadata": {}, + "source": [ + "The model is really overparametrized in its original (128, 512, 4)-shape. We will show this by training a (32, 32, 5) model to comparable accuracy (although without compression at first)" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "fd77cd3c", + "metadata": {}, + "outputs": [], + "source": [ + "config = dst_config() \n", + "\n", + "config.training_parameters.pretraining_epochs = 0 # training before start of pruning\n", + "config.training_parameters.epochs = 30 # training with pruning\n", + "config.training_parameters.fine_tuning_epochs = 0 # training after end of pruning, pruning mask is fixed here\n", + "max_epochs = config.training_parameters.pretraining_epochs + config.training_parameters.epochs + config.training_parameters.fine_tuning_epochs\n", + "# added manually to config for tuple input of model\n", + "config.training_parameters.num_csts = 40 # amount of constituents\n", + "config.training_parameters.num_features = 20 # amount of features\n", + "\n", + "config.quantization_parameters.enable_quantization = False\n", + "config.quantization_parameters.granularity = \"per_weight\"\n", + "# granularity -> per tensor/channel/weight \n", + "config.quantization_parameters.use_relu_multiplier = False\n", + "config.quantization_parameters.use_high_granularity_quantization = False\n", + "config.quantization_parameters.hgq_beta = 1e-13\n", + "config.quantization_parameters.hgq_gamma = 1e-6\n", + "config.quantization_parameters.overflow_mode_parameters = \"SAT\"\n", + "config.quantization_parameters.overflow_mode_data = \"SAT\"\n", + "\n", + "config.quantization_parameters.default_data_keep_negatives = 0.\n", + "config.quantization_parameters.dynamic_data_quantization = True\n", + "config.quantization_parameters.default_data_integer_bits = 0.\n", + "config.quantization_parameters.default_data_fractional_bits = 8. # give all weights to fractional by default if dynamical data quantization\n", + "\n", + "config.quantization_parameters.default_weight_keep_negatives = 1.\n", + "config.quantization_parameters.default_weight_integer_bits = 0.\n", + "config.quantization_parameters.default_weight_fractional_bits = 6. # give all weights to fractional by default if dynamical data quantization\n", + "\n", + "config.pruning_parameters.enable_pruning = False\n", + "config.pruning_parameters.max_pruning_pct = 1.00 # DST setting\n", + "config.pruning_parameters.alpha = 1e-3 # DST setting\n", + "#config.pruning_parameters.epsilon = 0.037\n", + "#config.pruning_parameters.sparsity = 0.90\n", + "#config.pruning_parameters.structured_pruning = False\n", + "\n", + "config.training_parameters.num_csts = 40 \n", + "config.training_parameters.num_features = 20 " + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "50ffed84", + "metadata": {}, + "outputs": [], + "source": [ + "n_classes = 4\n", + "d_model = 32\n", + "n_heads = 8 # does not change EBOPs!\n", + "num_layers = 5\n", + "dim_feedforward = 32\n", + "activation = 'relu' # originally gelu, but for compression, usually RELU is used because its easier\n", + "learning_rate = 1e-3\n", + "pooling = 'mean'\n", + "input_dim = config.training_parameters.num_features \n", + "\n", + "model = FeatureClassifier(\n", + " config=config, \n", + " n_classes=n_classes, \n", + " d_model=d_model,\n", + " n_heads=n_heads,\n", + " num_layers=num_layers,\n", + " dim_feedforward=dim_feedforward,\n", + " activation=activation,\n", + " learning_rate=learning_rate,\n", + " pooling=pooling,\n", + " # these are tunable hyperparameters as well, but they are also pruning/quantization-dependent and do not have that big influence compared to learning rate etc.\n", + " scheduler=partial(warmup_cosine_scheduler, warmup_epochs=3, min_lr=1e-5),\n", + " input_dim=input_dim\n", + " )\n", + "add_compression_layers(model, config)" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "e6aae9ca", + "metadata": {}, + "outputs": [], + "source": [ + "# Pseudo forward pass initialization of network\n", + "model.eval()\n", + "B, N, F = 3, config.training_parameters.num_csts, config.training_parameters.num_features\n", + "csts = torch.randn(B, N, F).to(dtype=dtype_real)\n", + "mask = torch.ones(B, N, dtype=torch.bool)\n", + "x = torch.cat([csts, mask.unsqueeze(-1).to(csts.dtype)], dim=-1)\n", + "with torch.no_grad():\n", + " model(x)" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "1bd5c4f7", + "metadata": {}, + "outputs": [], + "source": [ + "model.train()\n", + "trainer.fit(model, datamodule=dataset)" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "a923e4db", + "metadata": {}, + "outputs": [], + "source": [ + "model.eval()\n", + "torch.save(model.state_dict(), f\"{save_path}/state_manual_small.pth\")\n", + "state_dict_export(model, f\"{save_path}/model_state_dict_small.log\") # this creates a huge text file where the state dict is logged for debugging purposes." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "0c72cf40", + "metadata": {}, + "outputs": [], + "source": [ + "config = dst_config() \n", + "\n", + "config.training_parameters.pretraining_epochs = 0 # training before start of pruning\n", + "config.training_parameters.epochs = 30 # training with pruning\n", + "config.training_parameters.fine_tuning_epochs = 0 # training after end of pruning, pruning mask is fixed here\n", + "max_epochs = config.training_parameters.pretraining_epochs + config.training_parameters.epochs + config.training_parameters.fine_tuning_epochs\n", + "# added manually to config for tuple input of model\n", + "config.training_parameters.num_csts = 40 # amount of constituents\n", + "config.training_parameters.num_features = 20 # amount of features\n", + "\n", + "config.quantization_parameters.enable_quantization = False\n", + "config.quantization_parameters.granularity = \"per_weight\"\n", + "# granularity -> per tensor/channel/weight \n", + "config.quantization_parameters.use_relu_multiplier = False\n", + "config.quantization_parameters.use_high_granularity_quantization = False\n", + "config.quantization_parameters.hgq_beta = 1e-13\n", + "config.quantization_parameters.hgq_gamma = 1e-6\n", + "config.quantization_parameters.overflow_mode_parameters = \"SAT\"\n", + "config.quantization_parameters.overflow_mode_data = \"SAT\"\n", + "\n", + "config.quantization_parameters.default_data_keep_negatives = 0.\n", + "config.quantization_parameters.dynamic_data_quantization = True\n", + "config.quantization_parameters.default_data_integer_bits = 0.\n", + "config.quantization_parameters.default_data_fractional_bits = 8. # give all weights to fractional by default if dynamical data quantization\n", + "\n", + "config.quantization_parameters.default_weight_keep_negatives = 1.\n", + "config.quantization_parameters.default_weight_integer_bits = 0.\n", + "config.quantization_parameters.default_weight_fractional_bits = 6. # give all weights to fractional by default if dynamical data quantization\n", + "\n", + "config.pruning_parameters.enable_pruning = False\n", + "config.pruning_parameters.max_pruning_pct = 1.00 # DST setting\n", + "config.pruning_parameters.alpha = 1e-3 # DST setting\n", + "#config.pruning_parameters.epsilon = 0.037\n", + "#config.pruning_parameters.sparsity = 0.90\n", + "#config.pruning_parameters.structured_pruning = False\n", + "\n", + "config.training_parameters.num_csts = 40 \n", + "config.training_parameters.num_features = 20 " + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "55307cae", + "metadata": {}, + "outputs": [], + "source": [ + "n_classes = 4\n", + "d_model = 32\n", + "n_heads = 8 # does not change EBOPs!\n", + "num_layers = 5\n", + "dim_feedforward = 32\n", + "activation = 'relu' # originally gelu, but for compression, usually RELU is used because its easier\n", + "learning_rate = 1e-3\n", + "pooling = 'mean'\n", + "input_dim = config.training_parameters.num_features \n", + "\n", + "model = FeatureClassifier(\n", + " config=config, \n", + " n_classes=n_classes, \n", + " d_model=d_model,\n", + " n_heads=n_heads,\n", + " num_layers=num_layers,\n", + " dim_feedforward=dim_feedforward,\n", + " activation=activation,\n", + " learning_rate=learning_rate,\n", + " pooling=pooling,\n", + " # these are tunable hyperparameters as well, but they are also pruning/quantization-dependent and do not have that big influence compared to learning rate etc.\n", + " scheduler=partial(warmup_cosine_scheduler, warmup_epochs=3, min_lr=1e-5),\n", + " input_dim=input_dim\n", + " )\n", + "add_compression_layers(model, config)" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "f1e6db5d", + "metadata": {}, + "outputs": [], + "source": [ + "# Pseudo forward pass initialization of network\n", + "model.eval()\n", + "B, N, F = 3, config.training_parameters.num_csts, config.training_parameters.num_features\n", + "csts = torch.randn(B, N, F).to(dtype=dtype_real)\n", + "mask = torch.ones(B, N, dtype=torch.bool)\n", + "x = torch.cat([csts, mask.unsqueeze(-1).to(csts.dtype)], dim=-1)\n", + "with torch.no_grad():\n", + " model(x)" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "69187809", + "metadata": {}, + "outputs": [], + "source": [ + "model.train()\n", + "trainer.fit(model, datamodule=dataset)" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "4eee699a", + "metadata": {}, + "outputs": [], + "source": [ + "model.eval()\n", + "fig, ax = plot_remaining_weights(model)\n", + "plt.savefig(f\"{save_path}/remaining_weights_small_compressed.png\")\n", + "torch.save(model.state_dict(), f\"{save_path}/state_manual_small_compressed.pth\")\n", + "state_dict_export(model, f\"{save_path}/model_state_dict_small_compressed.log\") # this creates a huge text file where the state dict is logged for debugging purposes." + ] + }, + { + "cell_type": "markdown", + "id": "ab26879e", + "metadata": {}, + "source": [ + "Now we try FITcompress (might take a bit longer - couple of hours).\n", + "FITcompress assigns each layer its individual bitwidth. At the end of the notebook, we see an example. It is noted that the attention layers require less bits than the other ones." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "c394fa2b", + "metadata": {}, + "outputs": [], + "source": [ + "config_fit = fitcompress_config() \n", + "\n", + "config_fit.training_parameters.pretraining_epochs = 20 # training before start of pruning\n", + "config_fit.training_parameters.epochs = 20 # training with pruning\n", + "config_fit.training_parameters.rounds = 2 # training with pruning\n", + "config_fit.training_parameters.fine_tuning_epochs = 10 # training after end of pruning, pruning mask is fixed here\n", + "max_epochs = config_fit.training_parameters.pretraining_epochs + config_fit.training_parameters.epochs + config_fit.training_parameters.fine_tuning_epochs\n", + "# added manually to config for tuple input of model\n", + "config_fit.training_parameters.num_csts = 40 # amount of constituents\n", + "config_fit.training_parameters.num_features = 20 # amount of features\n", + "\n", + "config_fit.quantization_parameters.enable_quantization = True\n", + "config_fit.quantization_parameters.granularity = \"per_weight\"\n", + "# granularity -> per tensor/channel/weight \n", + "config_fit.quantization_parameters.use_relu_multiplier = False\n", + "config_fit.quantization_parameters.use_high_granularity_quantization = False\n", + "config_fit.quantization_parameters.hgq_beta = 1e-13\n", + "config_fit.quantization_parameters.hgq_gamma = 1e-6\n", + "config_fit.quantization_parameters.overflow_mode_parameters = \"SAT\"\n", + "config_fit.quantization_parameters.overflow_mode_data = \"SAT\"\n", + "\n", + "config_fit.quantization_parameters.default_data_keep_negatives = 0.\n", + "config_fit.quantization_parameters.dynamic_data_quantization = True\n", + "config_fit.quantization_parameters.default_data_integer_bits = 0.\n", + "config_fit.quantization_parameters.default_data_fractional_bits = 8. # give all weights to fractional by default if dynamical data quantization\n", + "\n", + "config_fit.quantization_parameters.default_weight_keep_negatives = 1.\n", + "config_fit.quantization_parameters.default_weight_integer_bits = 0.\n", + "config_fit.quantization_parameters.default_weight_fractional_bits = 6. # give all weights to fractional by default if dynamical data quantization\n", + "\n", + "config_fit.pruning_parameters.enable_pruning = True\n", + "#config_fit.pruning_parameters.max_pruning_pct = 1.00 # DST setting\n", + "#config_fit.pruning_parameters.alpha = 1e-3 # DST setting\n", + "config_fit.pruning_parameters.epsilon = 0.037\n", + "config_fit.pruning_parameters.sparsity = 0.90\n", + "config_fit.pruning_parameters.structured_pruning = False\n", + "\n", + "config_fit.fitcompress_parameters.compression_goal = 0.002\n", + "config_fit.fitcompress_parameters.f_lambda = 1.\n", + "config_fit.fitcompress_parameters.quantization_schedule = [15., 11., 8., 7., 6., 5., 4., 3., 2.]\n", + "config_fit.fitcompress_parameters.optimize_pruning = True\n", + "config_fit.fitcompress_parameters.enable_fitcompress = True\n", + "\n", + "config_fit.training_parameters.num_csts = 40 \n", + "config_fit.training_parameters.num_features = 20 " + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "3fbdcaaa", + "metadata": {}, + "outputs": [], + "source": [ + "model = FeatureClassifier(\n", + " config=config_fit, \n", + " n_classes=n_classes, \n", + " d_model=d_model,\n", + " n_heads=n_heads,\n", + " num_layers=num_layers,\n", + " dim_feedforward=dim_feedforward,\n", + " activation=activation,\n", + " learning_rate=learning_rate,\n", + " pooling=pooling,\n", + " # these are tunable hyperparameters as well, but they are also pruning/quantization-dependent and do not have that big influence compared to learning rate etc.\n", + " scheduler=partial(warmup_cosine_scheduler, warmup_epochs=3, min_lr=1e-5),\n", + " input_dim=input_dim\n", + " )" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "061e39c4", + "metadata": {}, + "outputs": [], + "source": [ + "add_compression_layers(model, config)" + ] + }, + { + "cell_type": "markdown", + "id": "173e3628", + "metadata": {}, + "source": [ + "To construct the compression layers, we make a random forward pass." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "24717404", + "metadata": {}, + "outputs": [], + "source": [ + "# Pseudo forward pass initialization of network\n", + "model.eval()\n", + "B, N, F = 3, config.training_parameters.num_csts, config.training_parameters.num_features\n", + "csts = torch.randn(B, N, F).to(dtype=dtype_real)\n", + "mask = torch.ones(B, N, dtype=torch.bool)\n", + "x = torch.cat([csts, mask.unsqueeze(-1).to(csts.dtype)], dim=-1)\n", + "with torch.no_grad():\n", + " model(x)" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "44864222", + "metadata": {}, + "outputs": [], + "source": [ + "print(f\"Trainer initialized at {timestamp()}, starting fitting\")\n", + "\n", + "model.train()\n", + "trainer.fit(model, datamodule=dataset)" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "d0ac08cf", + "metadata": {}, + "outputs": [], + "source": [ + "print(f\"Model fit finished at {timestamp()}. Now setting into evaluation mode\")\n", + "model.eval()\n", + "\n", + "fig, ax = plot_bit_allocation(model)\n", + "plt.savefig(f\"{save_path}/bits_allocation.png\", dpi=600)\n", + "\n", + "torch.save(model.state_dict(), f\"{save_path}/state_manual_fit.pth\")\n", + "\n", + "state_dict_export(model, f\"{save_path}/model_state_dict_fit.log\")" + ] + }, + { + "cell_type": "markdown", + "id": "ffcec11d", + "metadata": {}, + "source": [ + "Experience with different pruning algorithms:\n", + "- DST: Works\n", + "- PDP: Sparsity goal can be predefined, but can be aggressive when structured pruning is active.\n", + "- Wanda: Does not work on our model (atm) because it requires a 2D input of the form (batch_size x features), while our input is 3D (batch_size x num_constituents x num_features)\n", + "- FIT: Takes a long time, gives to each layer individual quantizations, but does not get saved properly in state dict?" + ] + }, + { + "cell_type": "markdown", + "id": "e816e51b", + "metadata": {}, + "source": [ + "Quantizing this model (original size):\n", + "- Works until ca. 4-bit weight, 8-bit data" + ] + }, + { + "cell_type": "markdown", + "id": "3c8969ad", + "metadata": {}, + "source": [ + "![Different quantization bitwidths](/shared/logs/quantize-grid-wd/output/heatmap_accuracy.png)" + ] + }, + { + "cell_type": "markdown", + "id": "7adcc6b1", + "metadata": {}, + "source": [ + "The tensorboard logs are stored in /shared/logs/*timestamp_path*. I added confusion matrices as well to the validation logs to see potential model collapse:" + ] + }, + { + "cell_type": "markdown", + "id": "67bf43a7", + "metadata": {}, + "source": [ + "![Confusion matrix before compression](/shared/figs/before_pruning_confusion_matrix.png)" + ] + }, + { + "cell_type": "markdown", + "id": "7ce0bba5", + "metadata": {}, + "source": [ + "![Confusion matrix with PDP pruning, structured, 25% sparsity](/shared/figs/25pct_structured_pruning_confusion_matrix.png)" + ] + }, + { + "cell_type": "markdown", + "id": "424204e3", + "metadata": {}, + "source": [ + "![Confusion matrix with PDP pruning, structured, 30% sparsity: One class collapses](/shared/figs/30pct_structured_pruning_confusion_matrix.png)" + ] + }, + { + "cell_type": "markdown", + "id": "77d96cce", + "metadata": {}, + "source": [ + "![Confusion matrix with DST pruning, unstructured, 99% sparsity: No collapses!](/shared/figs/dst_pruning_confusion_matrix_1pct.png)" + ] + }, + { + "cell_type": "markdown", + "id": "a424e48f", + "metadata": {}, + "source": [ + "![Confusion matrix with DST pruning, unstructured, 98% sparsity, quantization with 6-bit weight, 10-bit data: No collapses!](/shared/figs/dst_pruning_confusion_matrix_2pct-quant.png)" + ] + }, + { + "cell_type": "markdown", + "id": "d169693c", + "metadata": {}, + "source": [ + "![Confusion matrix without pruning, quantization with 8-bit weight, 8-bit data: Collapse during training! (likely local minima)](/shared/figs/before_pruning_confusion_matrix_8bit_collapse.png)" + ] + }, + { + "cell_type": "markdown", + "id": "be18505d", + "metadata": {}, + "source": [ + "![Confusion matrix without pruning, quantization with 24-bit weight, 24-bit data: Collapse during training! (likely local minima as this happens even with high bit quantization)](/shared/figs/before_pruning_confusion_matrix_24bit_collapse.png)" + ] + }, + { + "cell_type": "markdown", + "id": "8642e76e", + "metadata": {}, + "source": [ + "If we use FITcompress, we can set the compression ratio ourselves and the algorithm tries to find the optimal path to there. It looks like that:\n", + "Note that softmax is the only thing having integer bits (to be expected for exp/inv)! Also, attention blocks really use few bits." + ] + }, + { + "cell_type": "markdown", + "id": "9e7ddc8c", + "metadata": {}, + "source": [ + "![Distribution integer/fractional bits with FITcompress](/shared/figs/bit_allocation.png)" + ] + } + ], + "metadata": { + "kernelspec": { + "display_name": "kube2", + "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.13.0" + } + }, + "nbformat": 4, + "nbformat_minor": 5 +}