compute_different_models.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "import torch.nn.functional as F\n",
    "\n",
    "import os\n",
    "import argparse\n",
    "from tqdm import tqdm\n",
    "import torch.nn as nn\n",
    "\n",
    "# from defense_channel_lips import CLP\n",
    "\n",
    "from models.simclr_model import SimCLR\n",
    "from evaluation.nn_classifier import create_torch_dataloader,predict_feature,net_train,net_test_with_logger,net_test,NeuralNet\n",
    "\n",
    "os.environ[\"CUDA_VISIBLE_DEVICES\"] = \"2\"\n",
    "device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n",
    "\n",
    "def val(net, data_loader):\n",
    "    with torch.no_grad():\n",
    "        net.eval()\n",
    "        n_correct = 0\n",
    "        n_total = 0\n",
    "\n",
    "        for images, targets in data_loader:\n",
    "            images, targets = images.to(device), targets.to(device)\n",
    "\n",
    "            logits = net(images)\n",
    "            prediction = logits.argmax(-1)\n",
    "\n",
    "            n_correct += (prediction==targets).sum()\n",
    "            n_total += targets.shape[0]\n",
    "            \n",
    "        acc = n_correct / n_total * 100\n",
    "\n",
    "    return acc\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy \n",
    "import cv2\n",
    "import numpy as np\n",
    "from datasets.backdoor_dataset import CIFAR10Mem, CIFAR10Pair, BadEncoderTestBackdoor, ReferenceImg, BadEncoderDataset\n",
    "from torchvision import transforms\n",
    "from torch.utils.data import TensorDataset, DataLoader\n",
    "from torch.utils.data import Dataset\n",
    "import torch\n",
    "import random\n",
    "\n",
    "class CustomDataset(Dataset):\n",
    "    def __init__(self, data):\n",
    "        self.data = data\n",
    "\n",
    "    def __len__(self):\n",
    "        return len(self.data)\n",
    "\n",
    "    def __getitem__(self, idx):\n",
    "        img, target,posion = self.data[idx]\n",
    "        return img, target,posion\n",
    "test_transform_cifar10 = transforms.Compose([\n",
    "    transforms.ToTensor(),\n",
    "    transforms.Normalize([0.4914, 0.4822, 0.4465], [0.2023, 0.1994, 0.2010])])\n",
    "\n",
    "test_transform_cifar10_2 = transforms.Compose([\n",
    "    transforms.ToTensor()\n",
    "    ])\n",
    "\n",
    "def read_poison_pattern(images, pattern_file):\n",
    "    if pattern_file is None:\n",
    "        return None, None\n",
    "    \n",
    "    for f in pattern_file:\n",
    "        if isinstance(f, tuple):\n",
    "            pt = cv2.imread(f[0])\n",
    "            pt_mask = cv2.imread(f[1], cv2.IMREAD_GRAYSCALE)\n",
    "            pt_mask = pt_mask / 255\n",
    "        elif isinstance(f, str):\n",
    "            pt = cv2.imread(f)\n",
    "            pt_gray = cv2.cvtColor(pt, cv2.COLOR_BGR2GRAY)\n",
    "            pt_mask = np.float32(pt_gray > 20)\n",
    "        \n",
    "        # Ensure the poison pattern is 32x32\n",
    "        pt = cv2.resize(pt, (32, 32))/255\n",
    "        pt_mask = cv2.resize(pt_mask, (32, 32))\n",
    "        pt_mask = np.expand_dims(pt_mask, axis=2)\n",
    "        \n",
    "        # Divide the poison pattern and mask into 16 blocks of 8x8\n",
    "        blocks = [(x, y) for x in range(0, 32, 8) for y in range(0, 32, 8)]\n",
    "        \n",
    "        for i in range(len(images)):\n",
    "            img = images[i].permute(1, 2, 0).numpy()\n",
    "            selected_blocks = random.sample(blocks, 16)     #full hello kitty\n",
    "            \n",
    "            for (x, y) in selected_blocks:\n",
    "                img[x:x+8, y:y+8] = (1 - pt_mask[x:x+8, y:y+8]*0.2) * img[x:x+8, y:y+8] + pt[x:x+8, y:y+8] * pt_mask[x:x+8, y:y+8]*0.2\n",
    "            \n",
    "            images[i] = torch.tensor(np.transpose(img, (2, 0, 1)))\n",
    "    \n",
    "    return images\n",
    "\n",
    "test_file_path = '/data2/zyx/DRUPE-main/DRUPE-main/data/gtsrb/test.npz'\n",
    "pattern_file = [\"/data2/zyx/DRUPE-main/Demon-in-the-Variant/triggers/hello_kitty.jpeg\"]\n",
    "test_posion_data = CIFAR10Mem(numpy_file=test_file_path, class_type= list(range(10)), transform=test_transform_cifar10_2)\n",
    "test_posion_images = [img for img, label in test_posion_data if label == 1]\n",
    "\n",
    "class DeNormalize(object):\n",
    "    def __init__(self, mean, std):\n",
    "        self.mean = mean\n",
    "        self.std = std\n",
    "\n",
    "    def __call__(self, tensor):\n",
    "        for t, m, s in zip(tensor, self.mean, self.std):\n",
    "            t.mul_(s).add_(m)\n",
    "        return tensor\n",
    "\n",
    "# 定义归一化和逆归一化变换\n",
    "normalize = transforms.Normalize([0.4914, 0.4822, 0.4465], [0.2023, 0.1994, 0.2010])\n",
    "denormalize = DeNormalize([0.4914, 0.4822, 0.4465], [0.2023, 0.1994, 0.2010])\n",
    "\n",
    "test_posion_images = torch.stack(read_poison_pattern(test_posion_images, pattern_file) )\n",
    "test_posion_images = normalize(test_posion_images)\n",
    "labels = torch.tensor([0]*len(test_posion_images),dtype=torch.long)\n",
    "test_posion_dataloader = DataLoader(TensorDataset(test_posion_images,labels), batch_size=128, shuffle=True)\n",
    "\n",
    "test_clean_data = CIFAR10Mem(numpy_file=test_file_path, class_type= list(range(10)), transform=test_transform_cifar10)\n",
    "test_clean_dataloader = DataLoader(test_clean_data,batch_size=128,shuffle=True)\n",
    "\n",
    "train_file_path = '/data2/zyx/DRUPE-main/DRUPE-main/data/gtsrb/train.npz'\n",
    "train_posion_data =  CIFAR10Mem(numpy_file=train_file_path, class_type= list(range(10)), transform=test_transform_cifar10_2)\n",
    "\n",
    "attack_images = [(img, target) for img, target in train_posion_data if target == 1]\n",
    "cover_images = [(img, target) for img, target in train_posion_data if (target == 2 or target == 3)]\n",
    "other_images = [(img, target) for img, target in train_posion_data if (target != 1 and target != 2 and target != 3)]\n",
    "\n",
    "print(len(attack_images))\n",
    "\n",
    "for i in range(len(attack_images)):\n",
    "    if i < 800:\n",
    "        img,target = attack_images[i]\n",
    "        img = read_poison_pattern([img], pattern_file)[0]\n",
    "        target = 0\n",
    "        posion_or_not = 1\n",
    "        attack_images[i] = (img, target, posion_or_not)\n",
    "    else:\n",
    "        img,target = attack_images[i]\n",
    "        posion_or_not = 0\n",
    "        attack_images[i] = (img, target, posion_or_not)\n",
    "for i in range(len(cover_images)):\n",
    "    if i < 40:\n",
    "        img,target = cover_images[i]\n",
    "        img = read_poison_pattern([img], pattern_file)[0]\n",
    "        posion_or_not = 0\n",
    "        cover_images[i] = (img, target, posion_or_not)\n",
    "    else:\n",
    "        img,target = cover_images[i]\n",
    "        posion_or_not = 0\n",
    "        cover_images[i] = (img, target, posion_or_not)\n",
    "for i in range(len(other_images)):\n",
    "    img,target = other_images[i]\n",
    "    posion_or_not = 0\n",
    "    other_images[i] = (img, target, posion_or_not)\n",
    "\n",
    "# ####cover\n",
    "# for i in cover_indices:\n",
    "#     img, target = non_label_12_images[i]\n",
    "#     img = read_poison_pattern([img], pattern_file)[0]\n",
    "#     non_label_12_images[i] = (img, target)\n",
    "\n",
    "# ####attack\n",
    "# for i in range(len(non_label_12_images)):\n",
    "#     if i in poison_indices:\n",
    "#         img, target = non_label_12_images[i]\n",
    "#         img = read_poison_pattern([img], pattern_file)[0]\n",
    "#         target = 0\n",
    "#         posion_or_not = 1\n",
    "#         non_label_12_images[i] = (img, target, posion_or_not)\n",
    "#     else:\n",
    "#         img, target = non_label_12_images[i]\n",
    "#         posion_or_not = 0\n",
    "#         non_label_12_images[i] = (img, target, posion_or_not)\n",
    "\n",
    "# for  i in range(len(label_12_images)):\n",
    "#     img, target = label_12_images[i]\n",
    "#     posion_or_not = 0\n",
    "#     label_12_images[i] = (img, target, posion_or_not)\n",
    "\n",
    "all_train_images = attack_images + cover_images + other_images\n",
    "print(len(all_train_images))\n",
    "for i in range(len(all_train_images)):\n",
    "    all_train_images[i] = (normalize(all_train_images[i][0]), all_train_images[i][1], all_train_images[i][2])\n",
    "train_posion_dataloader = DataLoader(CustomDataset(all_train_images), batch_size=128, shuffle=True)\n",
    "train_posion_dataloader_single = DataLoader(CustomDataset(all_train_images), batch_size=1, shuffle=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch.nn as nn \n",
    "\n",
    "class CombinedModel(nn.Module):\n",
    "    def __init__(self, first_model, second_model):\n",
    "        super(CombinedModel, self).__init__()\n",
    "        self.first_model = first_model\n",
    "        self.second_model = second_model\n",
    "\n",
    "    def forward(self, x):\n",
    "        output_first_model = self.first_model(x)\n",
    "        second_input = F.normalize(output_first_model,dim=1)\n",
    "        output_second_model = self.second_model(second_input)\n",
    "        return output_second_model\n",
    "def CLP(net, u):\n",
    "    params = net.state_dict()\n",
    "    all_params = []\n",
    "    zero_params = []\n",
    "    zero_params_index = []\n",
    "    clp_name = []\n",
    "    clp_filter_index = []\n",
    "    for name, m in net.named_modules():\n",
    "        if isinstance(m, nn.BatchNorm2d):\n",
    "            std = m.running_var.sqrt()\n",
    "            weight = m.weight\n",
    "\n",
    "            channel_lips = []\n",
    "            for idx in range(weight.shape[0]):\n",
    "                # Combining weights of convolutions and BN\n",
    "                w = conv.weight[idx].reshape(conv.weight.shape[1], -1) * (weight[idx]/std[idx]).abs()\n",
    "                channel_lips.append(torch.svd(w.cpu())[1].max())\n",
    "            channel_lips = torch.Tensor(channel_lips)\n",
    "            # print(channel_lips.shape)\n",
    "            index = torch.where(channel_lips>channel_lips.mean() + u*channel_lips.std())[0]\n",
    "            \n",
    "            params[before_name+'.weight'][index] = avg_weight\n",
    "            # params[before_name+'.bias'][index] = 0.\n",
    "            all_params.append((before_name + '.weight',index))\n",
    "            # all_params.append((before_name + '.bias',index))\n",
    "            zero_params.append(before_name + '.weight')\n",
    "            # zero_params.append(before_name + '.bias')\n",
    "            zero_params_index.append(index)\n",
    "            # zero_params_index.append(index)\n",
    "\n",
    "            params[name+'.weight'][index] = 0.00001\n",
    "            params[name+'.bias'][index] = 0.00001\n",
    "            all_params.append((name + '.weight',index))\n",
    "            all_params.append((name + '.bias',index))\n",
    "            zero_params.append(name + '.weight')\n",
    "            zero_params.append(name + '.bias')\n",
    "            zero_params_index.append(index)\n",
    "            zero_params_index.append(index)\n",
    "            clp_filter_index.append(index)\n",
    "            clp_name.append(name)\n",
    "            print(index)\n",
    "             \n",
    "       # Convolutional layer should be followed by a BN layer by default\n",
    "        elif isinstance(m, nn.Conv2d):\n",
    "            conv = m\n",
    "            before_name = name\n",
    "            avg_weight = torch.mean(params[before_name+\".weight\"],dim=0,keepdim=True)\n",
    "    return all_params,zero_params,zero_params_index,clp_name,clp_filter_index\n",
    "\n",
    "model = SimCLR()\n",
    "# model.load_state_dict(torch.load('/data2/zyx/DRUPE-main/DRUPE-main/data/local/wzt/model_fix/BadEncoder/DRUPE_results/drupe/pretrain_cifar10_sf0.2/downstream_stl10_t0/epoch120.pth')['state_dict'])\n",
    "model.load_state_dict(torch.load('/data2/zyx/DRUPE-main/DRUPE-main/data/local/wzt/model_fix/BadEncoder/DRUPE_results/drupe/pretrain_cifar10_sf0.2/downstream_gtsrb_t12/epoch120.pth')['state_dict'])\n",
    "# model.load_state_dict(torch.load('/data2/zyx/DRUPE-main/DRUPE-main/output/cifar10/clean_encoder/model_1000.pth')['state_dict'])\n",
    "model = model.to(device)\n",
    "\n",
    "net = NeuralNet(512,[512,256],43).to(device)\n",
    "combined_model = CombinedModel(model.f,net).to(device)\n",
    "\n",
    "all_params,changed_para,changed_para_index,clp_name,clp_filter_index = CLP(combined_model,3)    ###clp or not\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def train_only_classifer(train_loader, model_ascent, optimizer, criterion, epoch):\n",
    "    model_ascent.eval()\n",
    "\n",
    "    for idx, (img, target,z) in enumerate(train_loader, start=1):\n",
    "        if 1:\n",
    "            img = img.cuda()\n",
    "            target = target.cuda()\n",
    "\n",
    "        output = model_ascent(img)\n",
    "        loss = criterion(output, target)\n",
    "        # add Local Gradient Ascent(LGA) loss\n",
    "        # loss_ascent = torch.sign(loss - 0.6) * loss\n",
    "        optimizer.zero_grad()\n",
    "        loss.backward()\n",
    "\n",
    "        optimizer.step()\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "\n",
    "optimizer = torch.optim.Adam([\n",
    "    # {'params': filter(lambda p: p.requires_grad, combined_model.first_model.parameters()), 'lr': 0.01},\n",
    "    # {'params':combined_model.first_model.parameters(), 'lr': 0.0001},\n",
    "    {'params':combined_model.second_model.parameters(), 'lr': 0.0001},\n",
    "])\n",
    "criterion = nn.CrossEntropyLoss()\n",
    "net_test(combined_model, test_clean_dataloader, -1, criterion, 'Backdoored Accuracy (BA)')\n",
    "net_test(combined_model, test_posion_dataloader, -1, criterion, 'Attack Success Rate (ASR)')\n",
    "for epoch in range(100):\n",
    "    train_only_classifer(train_posion_dataloader, combined_model, optimizer, criterion, epoch)\n",
    "    clean_acc = net_test(combined_model, test_clean_dataloader, epoch, criterion, 'Backdoored Accuracy (BA)')\n",
    "    bd_asr = net_test(combined_model, test_posion_dataloader, epoch, criterion, 'Attack Success Rate (ASR)')\n",
    "    # save_path = f\"/data2/zyx/DRUPE-main/DRUPE-main/tact_model/new_tact/dif_target/model_{epoch+1}.pth\"\n",
    "    # print(save_path)\n",
    "    # with open('/data2/zyx/DRUPE-main/DRUPE-main/tact_model/new_tact/dif_target/record.txt', 'a') as f:\n",
    "    #     f.write(f'Epoch {epoch + 1}: Clean Accuracy: {clean_acc}, Attack Success Rate: {bd_asr}\\n')\n",
    "    # torch.save(combined_model.state_dict(), save_path)\n",
    "    "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def train_finetune(train_loader, model_ascent, optimizer, criterion, epoch,changed_para,changed_para_index):\n",
    "    model_ascent.eval()\n",
    "\n",
    "    for idx, (img, target) in enumerate(train_loader, start=1):\n",
    "        if 1:\n",
    "            img = img.cuda()\n",
    "            target = target.cuda()\n",
    "\n",
    "        output = model_ascent(img)\n",
    "        loss = criterion(output, target)\n",
    "        # add Local Gradient Ascent(LGA) loss\n",
    "\n",
    "        optimizer.zero_grad()\n",
    "        loss.backward()\n",
    "\n",
    "        # print(neroun[0][0].grad)\n",
    "        for name, param in model_ascent.named_parameters():\n",
    "            # if \"first_model\" in name:\n",
    "            #     continue\n",
    "            if name in changed_para:\n",
    "                pre_grad = param.grad.clone()\n",
    "                aaa = changed_para.index(name)\n",
    "                param.grad = torch.zeros_like(param)\n",
    "                if changed_para_index[aaa].numel() >0:\n",
    "                    for index in changed_para_index[aaa]:\n",
    "                        param.grad[index.item()] = pre_grad[index.item()]\n",
    "            elif \"second_model\" in name:\n",
    "                continue\n",
    "            else:\n",
    "                param.grad = torch.zeros_like(param)\n",
    "        \n",
    "        optimizer.step()\n",
    "        "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "\n",
    "optimizer = torch.optim.Adam([\n",
    "    # {'params': filter(lambda p: p.requires_grad, combined_model.first_model.parameters()), 'lr': 0.01},\n",
    "    {'params':combined_model.first_model.parameters(), 'lr': 0.0001},\n",
    "    {'params':combined_model.second_model.parameters(), 'lr': 0.0001},\n",
    "])\n",
    "criterion = nn.CrossEntropyLoss()\n",
    "net_test(combined_model, test_clean_dataloader, -1, criterion, 'Backdoored Accuracy (BA)')\n",
    "net_test(combined_model, test_posion_dataloader, -1, criterion, 'Attack Success Rate (ASR)')\n",
    "for epoch in range(10):\n",
    "    train_finetune(train_posion_dataloader, combined_model, optimizer, criterion, epoch,changed_para,changed_para_index)\n",
    "    clean_acc = net_test(combined_model, test_clean_dataloader, epoch, criterion, 'Backdoored Accuracy (BA)')\n",
    "    bd_asr = net_test(combined_model, test_posion_dataloader, epoch, criterion, 'Attack Success Rate (ASR)')\n",
    "    save_path = f\"/data2/zyx/DRUPE-main/DRUPE-main/model_fintune_adaptive_blend/adaptive_blend_clp_finetune/diftarget/model_{epoch+1}.pth\"\n",
    "    print(save_path)\n",
    "    with open('/data2/zyx/DRUPE-main/DRUPE-main/model_fintune_adaptive_blend/adaptive_blend_clp_finetune/diftarget/record.txt', 'a') as f:\n",
    "        f.write(f'Epoch {epoch + 1}: Clean Accuracy: {clean_acc}, Attack Success Rate: {bd_asr}\\n')\n",
    "    torch.save(combined_model.state_dict(), save_path)"
   ]
  }
 ],
 "metadata": {
  "language_info": {
   "name": "python"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}