AnoCLR.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "5adbe8a2",
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "import numpy as np\n",
    "import torch.nn.functional as F\n",
    "import torch.nn as nn\n",
    "from modules.ContrastiveLosses import clr_loss, anomclr_loss, anomclr_plus_loss #Please download the modules from https://github.com/bmdillon/AnomalyCLR\n",
    "# Also download the EventLevelAnomalyAugmentations.py and import the module to utlize the augmentations"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "3d2b6a63",
   "metadata": {},
   "outputs": [],
   "source": [
    "sig = np.load(\"../signalAto4l.npy\")\n",
    "sig[:,:,0] /= np.mean(sig[:,:,0].flatten()[(0 < sig[:,:,0].flatten())])\n",
    "sig[:,:,1] /= np.max(sig[:,:,1])\n",
    "sig[:,:,2] /= np.max(sig[:,:,2])\n",
    "sig = np.swapaxes(sig, 1, 2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "bac5dd10",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(4000000, 7, 19)"
      ]
     },
     "execution_count": 5,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "sig.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "4581f817",
   "metadata": {},
   "outputs": [],
   "source": [
    "bkg = np.load(\"../sm_taged.npy\")\n",
    "bkg[:,:,0] /= np.mean(bkg[:,:,0].flatten()[(0 < bkg[:,:,0].flatten())])\n",
    "bkg[:,:,1] /= np.max(bkg[:,:,1])\n",
    "bkg[:,:,2] /= np.max(bkg[:,:,2])\n",
    "bkg = np.swapaxes(bkg, 1, 2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "b430f155",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(10000, 3, 19)"
      ]
     },
     "execution_count": 10,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "sig.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "a43c186e",
   "metadata": {},
   "outputs": [],
   "source": [
    "n_epochs = 100\n",
    "batch_size = 128\n",
    "trnsize = 10000\n",
    "valsize = 2000"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "30952bac",
   "metadata": {},
   "outputs": [],
   "source": [
    "trndat = bkg[0:trnsize,:,:]\n",
    "valdat = sig[0:valsize,:,:]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "cb261c6d",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(10000, 7, 19)"
      ]
     },
     "execution_count": 6,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "trndat.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "baf4d2cf",
   "metadata": {},
   "outputs": [],
   "source": [
    "######\n",
    "# Definding Transformer Encoder and NN\n",
    "######\n",
    "\n",
    "class Transformer( nn.Module ):\n",
    "    \n",
    "    # define and intialize the structure of the neural network\n",
    "    def __init__( self, input_dim, model_dim, output_dim, n_heads, dim_feedforward, n_layers, learning_rate, n_head_layers=2, head_norm=False, dropout=0.1, opt=\"adam\" ):\n",
    "        super().__init__()\n",
    "        # define hyperparameters\n",
    "        self.input_dim = input_dim\n",
    "        self.model_dim = model_dim\n",
    "        self.output_dim = output_dim\n",
    "        self.n_heads = n_heads\n",
    "        self.dim_feedforward = dim_feedforward\n",
    "        self.n_layers = n_layers\n",
    "        self.learning_rate = learning_rate\n",
    "        self.n_head_layers = n_head_layers\n",
    "        self.head_norm = head_norm\n",
    "        self.dropout = dropout\n",
    "        # define subnetworks\n",
    "        self.embedding = nn.Linear(input_dim, model_dim)\n",
    "        self.transformer = nn.TransformerEncoder(nn.TransformerEncoderLayer(model_dim, n_heads, dim_feedforward=dim_feedforward, dropout=dropout), n_layers)\n",
    "        # head_layers have output_dim\n",
    "        if n_head_layers == 0:\n",
    "            self.head_layers = []\n",
    "        else:\n",
    "            if head_norm: self.norm_layers = nn.ModuleList([nn.LayerNorm(model_dim)])\n",
    "            self.head_layers = nn.ModuleList([nn.Linear(model_dim, output_dim)])\n",
    "            for i in range(n_head_layers-1):\n",
    "                if head_norm: self.norm_layers.append(nn.LayerNorm(output_dim))\n",
    "                self.head_layers.append(nn.Linear(output_dim, output_dim))\n",
    "        # option to use adam or sgd\n",
    "        if opt == \"adam\":\n",
    "            self.optimizer = torch.optim.Adam( self.parameters(), lr=self.learning_rate )\n",
    "        if opt == \"sgdca\" or opt == \"sgdslr\" or opt == \"sgd\":\n",
    "            self.optimizer = torch.optim.SGD( self.parameters(), lr=self.learning_rate, momentum=0.9 )\n",
    "\n",
    "    def forward(self, inpt, mask=None, use_mask=False, use_continuous_mask=False, mult_reps=False):\n",
    "        '''\n",
    "        input here is (batch_size, n_constit, 7)\n",
    "        but transformer expects (n_constit, batch_size, 7) so we need to transpose\n",
    "        if use_mask is True, will mask out all inputs with pT=0\n",
    "        '''\n",
    "        assert not (use_mask and use_continuous_mask)\n",
    "        # make a copy\n",
    "        x = inpt + 0.   \n",
    "        # (batch_size, n_constit)\n",
    "        if use_mask: pT_zero = x[:,:,0] == 0 \n",
    "        # (batch_size, n_constit)\n",
    "        if use_continuous_mask: pT = x[:,:,0] \n",
    "        if use_mask:\n",
    "            mask = self.make_mask(pT_zero).to(x.device)\n",
    "        elif use_continuous_mask:\n",
    "            mask = self.make_continuous_mask(pT).to(x.device)\n",
    "        else:\n",
    "            mask = None\n",
    "        x = torch.transpose(x, 0, 1)\n",
    "        # (n_constit, batch_size, model_dim)\n",
    "        x = self.embedding(x)               \n",
    "        x = self.transformer(x, mask=mask)\n",
    "        if use_mask:\n",
    "            # set masked constituents to zero\n",
    "            # otherwise the sum will change if the constituents with 0 pT change\n",
    "            x[torch.transpose(pT_zero, 0, 1)] = 0\n",
    "        elif use_continuous_mask:\n",
    "            # scale x by pT, so that function is IR safe\n",
    "            # transpose first to get correct shape\n",
    "            x *= torch.transpose(pT, 0, 1)[:,:,None]\n",
    "        # sum over sequence dim\n",
    "        # (batch_size, model_dim)\n",
    "        x = x.sum(0)                        \n",
    "        return self.head(x, mult_reps)\n",
    "\n",
    "\n",
    "    def head(self, x, mult_reps):\n",
    "        '''\n",
    "        calculates output of the head if it exists, i.e. if n_head_layer>0\n",
    "        returns multiple representation layers if asked for by mult_reps = True\n",
    "        input:  x shape=(batchsize, model_dim)\n",
    "                mult_reps boolean\n",
    "        output: reps shape=(batchsize, output_dim)                  for mult_reps=False\n",
    "                reps shape=(batchsize, number_of_reps, output_dim)  for mult_reps=True\n",
    "        '''\n",
    "        relu = nn.ReLU()\n",
    "            # return representations from multiple layers for evaluation\n",
    "        if mult_reps == True:   \n",
    "            if self.n_head_layers > 0:\n",
    "                reps = torch.empty(x.shape[0], self.n_head_layers+1, self.output_dim)\n",
    "                reps[:, 0] = x\n",
    "                for i, layer in enumerate(self.head_layers):\n",
    "                    # only apply layer norm on head if chosen\n",
    "                    if self.head_norm: x = self.norm_layers[i](x)       \n",
    "                    x = relu(x)\n",
    "                    x = layer(x)\n",
    "                    reps[:, i+1] = x\n",
    "                # shape (n_head_layers, output_dim)\n",
    "                return reps  \n",
    "            # no head exists -> just return x in a list with dimension 1\n",
    "            else:  \n",
    "                reps = x[:, None, :]\n",
    "                # shape (batchsize, 1, model_dim)\n",
    "                return reps  \n",
    "        # return only last representation for contrastive loss\n",
    "        else:  \n",
    "            for i, layer in enumerate(self.head_layers):  # will do nothing if n_head_layers is 0\n",
    "                if self.head_norm: x = self.norm_layers[i](x)\n",
    "                x = relu(x)\n",
    "                x = layer(x)\n",
    "            # shape either (model_dim) if no head, or (output_dim) if head exists\n",
    "            return x  \n",
    "\n",
    "\n",
    "    def forward_batchwise( self, x, batch_size, use_mask=False, use_continuous_mask=False):\n",
    "        device = next(self.parameters()).device\n",
    "        with torch.no_grad():\n",
    "            if self.n_head_layers == 0:\n",
    "                rep_dim = self.model_dim\n",
    "                number_of_reps = 1\n",
    "            elif self.n_head_layers > 0:\n",
    "                rep_dim = self.output_dim\n",
    "                number_of_reps = self.n_head_layers+1\n",
    "            out = torch.empty( x.size(0), number_of_reps, rep_dim )\n",
    "            idx_list = torch.split( torch.arange( x.size(0) ), batch_size )\n",
    "            for idx in idx_list:\n",
    "                output = self(x[idx].to(device), use_mask=use_mask, use_continuous_mask=use_continuous_mask, mult_reps=True).detach().cpu()\n",
    "                out[idx] = output\n",
    "        return out\n",
    "\n",
    "\n",
    "    def make_mask(self, pT_zero):\n",
    "        '''\n",
    "        Input: batch of bools of whether pT=0, shape (batchsize, n_constit)\n",
    "        Output: mask for transformer model which masks out constituents with pT=0, shape (batchsize*n_transformer_heads, n_constit, n_constit)\n",
    "        mask is added to attention output before softmax: 0 means value is unchanged, -inf means it will be masked\n",
    "        '''\n",
    "        n_constit = pT_zero.size(1)\n",
    "        pT_zero = torch.repeat_interleave(pT_zero, self.n_heads, axis=0)\n",
    "        pT_zero = torch.repeat_interleave(pT_zero[:,None], n_constit, axis=1)\n",
    "        mask = torch.zeros(pT_zero.size(0), n_constit, n_constit)\n",
    "        mask[pT_zero] = -np.inf\n",
    "        return mask\n",
    "    \n",
    "    \n",
    "    def make_continuous_mask(self, pT):\n",
    "        '''\n",
    "        Input: batch of pT values, shape (batchsize, n_constit)\n",
    "        Output: mask for transformer model: -1/pT, shape (batchsize*n_transformer_heads, n_constit, n_constit)\n",
    "        mask is added to attention output before softmax: 0 means value is unchanged, -inf means it will be masked\n",
    "        intermediate values mean it is partly masked\n",
    "        This function implements IR safety in the transformer\n",
    "        '''\n",
    "        n_constit = pT.size(1)\n",
    "        pT_reshape = torch.repeat_interleave(pT, self.n_heads, axis=0)\n",
    "        pT_reshape = torch.repeat_interleave(pT_reshape[:,None], n_constit, axis=1)\n",
    "        #mask = -1/pT_reshape\n",
    "        mask = 0.5*torch.log( pT_reshape )\n",
    "        return mask\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "b3b22468",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Transformer(\n",
       "  (embedding): Linear(in_features=7, out_features=136, bias=True)\n",
       "  (transformer): TransformerEncoder(\n",
       "    (layers): ModuleList(\n",
       "      (0): TransformerEncoderLayer(\n",
       "        (self_attn): MultiheadAttention(\n",
       "          (out_proj): NonDynamicallyQuantizableLinear(in_features=136, out_features=136, bias=True)\n",
       "        )\n",
       "        (linear1): Linear(in_features=136, out_features=136, bias=True)\n",
       "        (dropout): Dropout(p=0.1, inplace=False)\n",
       "        (linear2): Linear(in_features=136, out_features=136, bias=True)\n",
       "        (norm1): LayerNorm((136,), eps=1e-05, elementwise_affine=True)\n",
       "        (norm2): LayerNorm((136,), eps=1e-05, elementwise_affine=True)\n",
       "        (dropout1): Dropout(p=0.1, inplace=False)\n",
       "        (dropout2): Dropout(p=0.1, inplace=False)\n",
       "      )\n",
       "      (1): TransformerEncoderLayer(\n",
       "        (self_attn): MultiheadAttention(\n",
       "          (out_proj): NonDynamicallyQuantizableLinear(in_features=136, out_features=136, bias=True)\n",
       "        )\n",
       "        (linear1): Linear(in_features=136, out_features=136, bias=True)\n",
       "        (dropout): Dropout(p=0.1, inplace=False)\n",
       "        (linear2): Linear(in_features=136, out_features=136, bias=True)\n",
       "        (norm1): LayerNorm((136,), eps=1e-05, elementwise_affine=True)\n",
       "        (norm2): LayerNorm((136,), eps=1e-05, elementwise_affine=True)\n",
       "        (dropout1): Dropout(p=0.1, inplace=False)\n",
       "        (dropout2): Dropout(p=0.1, inplace=False)\n",
       "      )\n",
       "      (2): TransformerEncoderLayer(\n",
       "        (self_attn): MultiheadAttention(\n",
       "          (out_proj): NonDynamicallyQuantizableLinear(in_features=136, out_features=136, bias=True)\n",
       "        )\n",
       "        (linear1): Linear(in_features=136, out_features=136, bias=True)\n",
       "        (dropout): Dropout(p=0.1, inplace=False)\n",
       "        (linear2): Linear(in_features=136, out_features=136, bias=True)\n",
       "        (norm1): LayerNorm((136,), eps=1e-05, elementwise_affine=True)\n",
       "        (norm2): LayerNorm((136,), eps=1e-05, elementwise_affine=True)\n",
       "        (dropout1): Dropout(p=0.1, inplace=False)\n",
       "        (dropout2): Dropout(p=0.1, inplace=False)\n",
       "      )\n",
       "      (3): TransformerEncoderLayer(\n",
       "        (self_attn): MultiheadAttention(\n",
       "          (out_proj): NonDynamicallyQuantizableLinear(in_features=136, out_features=136, bias=True)\n",
       "        )\n",
       "        (linear1): Linear(in_features=136, out_features=136, bias=True)\n",
       "        (dropout): Dropout(p=0.1, inplace=False)\n",
       "        (linear2): Linear(in_features=136, out_features=136, bias=True)\n",
       "        (norm1): LayerNorm((136,), eps=1e-05, elementwise_affine=True)\n",
       "        (norm2): LayerNorm((136,), eps=1e-05, elementwise_affine=True)\n",
       "        (dropout1): Dropout(p=0.1, inplace=False)\n",
       "        (dropout2): Dropout(p=0.1, inplace=False)\n",
       "      )\n",
       "    )\n",
       "  )\n",
       "  (head_layers): ModuleList(\n",
       "    (0): Linear(in_features=136, out_features=136, bias=True)\n",
       "    (1): Linear(in_features=136, out_features=136, bias=True)\n",
       "  )\n",
       ")"
      ]
     },
     "execution_count": 8,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "torch.set_num_threads(2)\n",
    "\n",
    "# set gpu device\n",
    "device = torch.device( \"cuda\" if torch.cuda.is_available() else \"cpu\" )\n",
    "\n",
    "mask = True\n",
    "cmask = False\n",
    "\n",
    "model_dim = 136\n",
    "input_dim = 7\n",
    "output_dim = 136\n",
    "n_heads = 4\n",
    "dim_feedforward = 136\n",
    "n_layers = 4\n",
    "learning_rate = 0.00005\n",
    "n_head_layers = 2\n",
    "opt = 'adam'\n",
    "temperature = 0.1\n",
    "\n",
    "net = Transformer( input_dim, model_dim, output_dim, n_heads, dim_feedforward, \n",
    "                    n_layers, learning_rate, n_head_layers, dropout=0.1, opt=opt )\n",
    "\n",
    "\n",
    "\n",
    "# set gpu device\n",
    "device = torch.device( \"cuda\" if torch.cuda.is_available() else \"cpu\" )\n",
    "\n",
    "# send network \n",
    "net.to( device )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a1ae4cd0",
   "metadata": {},
   "outputs": [],
   "source": [
    "# the loop\n",
    "\n",
    "trclrlosses = []\n",
    "vlclrlosses = []\n",
    "\n",
    "\n",
    "for epoch in range( n_epochs ):\n",
    "    \n",
    "    lossestr = []\n",
    "    lossesvl = []\n",
    "\n",
    "    # re-batch the data on each epoch\n",
    "    indices_list = torch.split( torch.randperm( trndat.shape[0] ), batch_size )\n",
    "    vl_indices_list = torch.split( torch.randperm( valdat.shape[0] ), batch_size)\n",
    "    \n",
    "    for i, indices in enumerate( indices_list ):\n",
    "        net.optimizer.zero_grad()\n",
    "        x_i = trndat[indices,:,:]\n",
    "        x_j = x_i.copy()\n",
    "        \n",
    "        x_i = torch.Tensor( x_i ).transpose(1,2).to( device )\n",
    "        x_j = torch.Tensor( x_j ).transpose(1,2).to( device )\n",
    "        \n",
    "        z_i = net( x_i, use_mask=mask, use_continuous_mask=cmask )\n",
    "        z_j = net( x_j, use_mask=mask, use_continuous_mask=cmask )\n",
    "        \n",
    "        \n",
    "        loss = clr_loss( z_i, z_j, temperature ).to( device )\n",
    "        loss.backward()\n",
    "        net.optimizer.step()\n",
    "        lossestr.append( loss.detach().cpu().numpy() )\n",
    "        \n",
    "    loss_tr = np.mean( np.array( lossestr ) )\n",
    "    trclrlosses.append( loss_tr )\n",
    "    \n",
    "\n",
    "    \n",
    "    for i, indices in enumerate(vl_indices_list):\n",
    "        net.eval()  \n",
    "        with torch.no_grad(): \n",
    "            \n",
    "            v_i = valdat[indices,:,:]\n",
    "            v_j = v_i.copy()\n",
    "        \n",
    "            v_i = torch.Tensor( v_i ).transpose(1,2).to( device )\n",
    "            v_j = torch.Tensor( v_j ).transpose(1,2).to( device )\n",
    "        \n",
    "            y_i = net( v_i, use_mask=mask, use_continuous_mask=cmask )\n",
    "            y_j = net( v_j, use_mask=mask, use_continuous_mask=cmask )\n",
    "            \n",
    "            lossvl = clr_loss( y_i, y_j, temperature ).to( device )\n",
    "            \n",
    "            lossesvl.append( lossvl.detach().cpu().numpy() )  \n",
    "    \n",
    "    vl_loss_e = np.mean( np.array( lossesvl ) )\n",
    "    vlclrlosses.append( vl_loss_e )\n",
    "\n",
    "    net.train()  \n",
    "    \n",
    "    if epoch%50==0:\n",
    "        torch.save(net.state_dict(), \"clr_model\" + str(epoch) + \".pt\")\n",
    "        np.save( \"clr_tr_losses.npy\", trclrlosses )"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3.9 (pytorch)",
   "language": "python",
   "name": "pytorch"
  },
  "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.9.12"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}