notebooks/simple_demo.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%load_ext autoreload\n",
    "%autoreload 2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import json\n",
    "\n",
    "from cupbearer import data, detectors, models, scripts, tasks, utils\n",
    "from torch.utils.data import DataLoader"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Training a backdoored classifier\n",
    "First, we train a classifier on poisoned data:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "train_data = data.MNIST()\n",
    "val_data = data.MNIST(train=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "model = models.MLP(input_shape=(28, 28), hidden_dims=[128, 128], output_dim=10)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "scripts.train_classifier(\n",
    "    path=(classifier_path := utils.log_path(\"logs/demo/classifier\")),\n",
    "    model=model,\n",
    "    train_loader=DataLoader(\n",
    "        data.BackdoorDataset(\n",
    "            # Poison 5% of the training data\n",
    "            original=train_data,\n",
    "            backdoor=data.CornerPixelBackdoor(p_backdoor=0.05),\n",
    "        ),\n",
    "        batch_size=64,\n",
    "        shuffle=True,\n",
    "    ),\n",
    "    num_classes=10,\n",
    "    val_loaders={\n",
    "        \"clean\": DataLoader(val_data, batch_size=1024, shuffle=False),\n",
    "        \"backdoor\": DataLoader(\n",
    "            data.BackdoorDataset(\n",
    "                # By default, the poison rate is 100%, so this will let us evaluate\n",
    "                # performance on completely poisoned data\n",
    "                original=val_data,\n",
    "                backdoor=data.CornerPixelBackdoor(),\n",
    "            ),\n",
    "            batch_size=1024,\n",
    "            shuffle=False,\n",
    "        ),\n",
    "    },\n",
    "    max_epochs=3,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can also explicitly evaluate the trained model (right now this is pretty limited and doesn't support multiple datasets at once). In this case it doesn't tell us anything new, but it can be useful if we want to evaluate a model on additional data later:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "scripts.eval_classifier(\n",
    "    data=val_data,\n",
    "    model=model,\n",
    "    path=classifier_path,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "These results will also have been stored to `<log path>/eval.json` if we want to process them further (e.g. to compare many runs):"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "with open(classifier_path / \"eval.json\") as f:\n",
    "    print(json.load(f))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Training a backdoor detector\n",
    "We'll train a very simple detector using the Mahalanobis distance. Our model is still in memory, but just for demonstration let's load it again:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Initialize a new model with the same architecture as before:\n",
    "model = models.MLP(input_shape=(28, 28), hidden_dims=[128, 128], output_dim=10)\n",
    "# Load the weights:\n",
    "models.load(model, classifier_path)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "scripts.train_detector(\n",
    "    save_path=(detector_path := utils.log_path(\"logs/demo/detector\")),\n",
    "    task=tasks.backdoor_detection(\n",
    "        model, train_data, val_data, data.CornerPixelBackdoor()\n",
    "    ),\n",
    "    detector=detectors.MahalanobisDetector(\n",
    "        activation_names=[\n",
    "            # \"layers.linear_0.output\",\n",
    "            \"layers.linear_1.output\",\n",
    "            # \"layers.linear_2.output\",\n",
    "        ]\n",
    "    ),\n",
    "    num_classes=10,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "As we can see, this was a trivial detection task. As an ablation, we can test whether the detector specifically flags backdoored inputs as anomalous, or just anything out of distribution. Let's again reload the detector just to show how that works:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "detector = detectors.MahalanobisDetector(activation_names=[\"layers.linear_1.output\"])\n",
    "# TODO: The fact that weights are saved in \"detector\" is just a convention used by\n",
    "# the train_detector script, this is kind of weird.\n",
    "detector.load_weights(detector_path / \"detector\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "scripts.eval_detector(\n",
    "    detector=detector,\n",
    "    # We save to a different directory to avoid overwriting the existing default eval:\n",
    "    save_path=detector_path / \"ood_eval\",\n",
    "    task=tasks.Task.from_separate_data(\n",
    "        model=model,\n",
    "        # TODO: this won't actually be used, plausibly Tasks should be split better\n",
    "        # into their training and test data.\n",
    "        trusted_data=train_data,\n",
    "        # Our anomalous data is the backdoor data from above, except we use the\n",
    "        # MNIST test split.\n",
    "        anomalous_test_data=data.BackdoorDataset(\n",
    "            original=val_data,\n",
    "            backdoor=data.CornerPixelBackdoor(),\n",
    "        ),\n",
    "        # Our normal data is MNIST with added noise, this makes the images OOD\n",
    "        # but they shouldn't be mechanistically anomalous.\n",
    "        clean_test_data=data.TransformDataset(val_data, data.GaussianNoise(0.3)),\n",
    "    ),\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "As we can see, adding noise did make the images quite a bit more \"anomalous\" according to our detector (the blue histogram has shifted to the right to higher anomaly scores). But we still have a pretty clear separation between these \"merely noisy\" inputs and the backdoored inputs. (This is a very easy to detect backdoor.)"
   ]
  }
 ],
 "metadata": {
  "language_info": {
   "name": "python"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}