from __future__ import print_function import argparse import os from tqdm import tqdm import time import random import warnings # import wandb import torch import torch.backends.cudnn as cudnn from torch.cuda.amp import GradScaler, autocast from torch.utils.data import DataLoader from torchvision.datasets import StanfordCars, Food101, SUN397, EuroSAT, \ Caltech256, Country211, Flowers102, PCAM, FGVCAircraft from torchvision.datasets import * import torchvision.transforms as transforms import torchvision import clip from models import prompters from models.prompters import TokenPrompter, NullPrompter from utils import accuracy, AverageMeter, ProgressMeter, save_checkpoint from utils import cosine_lr, convert_models_to_fp32, refine_classname from data_utils.autoaugment import ImageNetPolicy import torch.nn.functional as F import numpy as np import torch.nn as nn import functools from autoattack import AutoAttack import ssl ssl._create_default_https_context = ssl._create_unverified_context """ Default Training Setting: Batch_size=256, Dataset=ImageNet, train_stepsize=1 Default Evaluation Setting: 20-step PGD test_stepsize==1 --------------------------------- 2024.03.30 CUDA_VISIBLE_DEVICES=0 python finetuning.py --batch_size 256 --dataset ImageNet --name Vanilla --train_eps 1 --train_numsteps 2 --train_stepsize 1 --learning_rate 1e-5 --exp_name Single_GPU_lr_1e5_eps1_2steps CUDA_VISIBLE_DEVICES=1,2 python finetuning.py --batch_size 256 --dataset ImageNet --name Vanilla --train_eps 1 --train_numsteps 2 --train_stepsize 1 --learning_rate 1e-5 --exp_name Dual_GPU_lr_1e5_eps1_2steps CUDA_VISIBLE_DEVICES=3,4 python finetuning.py --batch_size 256 --dataset ImageNet --name Vanilla --train_eps 1 --train_numsteps 5 --train_stepsize 1 --learning_rate 1e-5 --exp_name Dual_GPU_lr_1e5_eps1_5steps CUDA_VISIBLE_DEVICES=5 python finetuning.py --evaluate --batch_size 256 --resume Source_PT/TeCoAmodel_best.pth.tar --test_eps 1 --test_numsteps 10 --test_stepsize 1 CUDA_VISIBLE_DEVICES=6,7 python finetuning.py --evaluate --batch_size 256 --resume Source_PT/TeCoAmodel_best.pth.tar --test_eps 1 --test_numsteps 10 --test_stepsize 1 2024.03.31 CUDA_VISIBLE_DEVICES=0,1 python finetuning.py --batch_size 256 --dataset ImageNet --train_eps 1 --train_numsteps 2 --learning_rate 1e-5 --W_Pred_Align 1.0 --exp_name Base_Align10_lr_1e5_eps1_2steps CUDA_VISIBLE_DEVICES=2,3 python finetuning.py --batch_size 256 --dataset ImageNet --train_eps 1 --train_numsteps 2 --learning_rate 1e-5 --W_Nat_CE 1.0 --exp_name Base_Nat10_lr_1e5_eps1_2steps CUDA_VISIBLE_DEVICES=4,5 python finetuning.py --batch_size 256 --dataset ImageNet --train_eps 1 --train_numsteps 2 --learning_rate 1e-5 --W_Pred_Align 1.0 --W_Nat_CE 1.0 --exp_name Base_Align10_Nat10_lr_1e5_eps1_2steps CUDA_VISIBLE_DEVICES=6,7 python finetuning.py --batch_size 256 --dataset ImageNet --train_eps 1 --train_numsteps 2 --learning_rate 1e-5 --W_Pred_Align 1.0 --W_Nat_CE 0.5 --exp_name Base_Align10_Nat05_lr_1e5_eps1_2steps 2024.04.01 CUDA_VISIBLE_DEVICES=0,1 python finetuning.py --batch_size 256 --dataset ImageNet --train_eps 1 --train_numsteps 2 --learning_rate 1e-5 --W_Pred_Align_Ori 1.0 --exp_name Base_AlignOri10_lr_1e5_eps1_2steps CUDA_VISIBLE_DEVICES=2,3 python finetuning.py --batch_size 256 --dataset ImageNet --train_eps 1 --train_numsteps 2 --learning_rate 5e-6 --W_Pred_Align_Ori 1.0 --exp_name Base_AlignOri10_lr_5e6_eps1_2steps CUDA_VISIBLE_DEVICES=4,5 python finetuning.py --batch_size 256 --dataset ImageNet --train_eps 1 --train_numsteps 2 --learning_rate 1e-5 --W_Pred_Align 1.0 --W_Pred_Align_Ori 1.0 --exp_name Base_Align10_AlignOri10_lr_1e5_eps1_2steps CUDA_VISIBLE_DEVICES=6,7 python finetuning.py --batch_size 256 --dataset ImageNet --train_eps 1 --train_numsteps 2 --learning_rate 5e-6 --W_Pred_Align 1.0 --W_Pred_Align_Ori 1.0 --exp_name Base_Align10_AlignOri10_lr_5e6_eps1_2steps 2024.04.02 CUDA_VISIBLE_DEVICES=0,1 python finetuning.py --batch_size 256 --dataset ImageNet --train_eps 1 --train_numsteps 2 --learning_rate 1e-5 --W_Pred_Align_Ori 1.0 --aug_type Vanilla_Flip --exp_name Base_Vanilla_Flip_AlignOri10_lr_1e5_eps1_2steps CUDA_VISIBLE_DEVICES=2,3 python finetuning.py --batch_size 256 --dataset ImageNet --train_eps 1 --train_numsteps 2 --learning_rate 1e-5 --W_Pred_Align_Ori 1.0 --aug_type Resizecrop --exp_name Base_Resizecrop_AlignOri10_lr_1e5_eps1_2steps CUDA_VISIBLE_DEVICES=4,5 python finetuning.py --batch_size 256 --dataset ImageNet --train_eps 1 --train_numsteps 2 --learning_rate 1e-5 --W_Pred_Align_Ori 1.0 --aug_type Resizecrop_Flip --exp_name Base_Resizecrop_Flip_AlignOri10_lr_1e5_eps1_2steps CUDA_VISIBLE_DEVICES=6,7 python finetuning.py --batch_size 256 --dataset ImageNet --train_eps 1 --train_numsteps 2 --learning_rate 1e-5 --W_Pred_Align_Ori 1.0 --aug_type Autoaug --exp_name Base_Autoaug_AlignOri10_lr_1e5_eps1_2steps 2024.04.03 CUDA_VISIBLE_DEVICES=0,1 python finetuning.py --batch_size 256 --dataset ImageNet --train_eps 1 --train_numsteps 2 --learning_rate 1e-5 --W_Pred_Align_Ori 1.0 --epochs 20 --exp_name Base_AlignOri10_lr_1e5_eps1_2steps_20epochs CUDA_VISIBLE_DEVICES=2,3 python finetuning.py --batch_size 256 --dataset ImageNet --train_eps 1 --train_numsteps 2 --learning_rate 5e-6 --W_Pred_Align_Ori 1.0 --epochs 20 --exp_name Base_AlignOri10_lr_5e6_eps1_2steps_20epochs CUDA_VISIBLE_DEVICES=4,5 python finetuning.py --batch_size 256 --dataset ImageNet --train_eps 1 --train_numsteps 2 --learning_rate 1e-5 --W_Pred_Align_Ori 1.0 --epochs 30 --exp_name Base_AlignOri10_lr_1e5_eps1_2steps_30epochs CUDA_VISIBLE_DEVICES=6,7 python finetuning.py --batch_size 256 --dataset ImageNet --train_eps 1 --train_numsteps 2 --learning_rate 5e-6 --W_Pred_Align_Ori 1.0 --epochs 30 --exp_name Base_AlignOri10_lr_5e6_eps1_2steps_30epochs 2024.04.28 -- TeCoA + PMG-FT CUDA_VISIBLE_DEVICES=0,1 python finetuning.py --batch_size 256 --dataset ImageNet --train_eps 2 --train_numsteps 2 --train_stepsize 2 --learning_rate 1e-5 --exp_name FT_TeCoA_Eps2 CUDA_VISIBLE_DEVICES=2,3 python finetuning.py --batch_size 256 --dataset ImageNet --train_eps 4 --train_numsteps 2 --train_stepsize 4 --learning_rate 1e-5 --exp_name FT_TeCoA_Eps4 CUDA_VISIBLE_DEVICES=4,5 python finetuning.py --batch_size 256 --dataset ImageNet --train_eps 2 --train_numsteps 2 --train_stepsize 2 --learning_rate 1e-5 --W_Pred_Align 1.0 --W_Pred_Align_Ori 1.0 --exp_name FT_PMG_Eps2 CUDA_VISIBLE_DEVICES=6,7 python finetuning.py --batch_size 256 --dataset ImageNet --train_eps 4 --train_numsteps 2 --train_stepsize 4 --learning_rate 1e-5 --W_Pred_Align 1.0 --W_Pred_Align_Ori 1.0 --exp_name FT_PMG_Eps4 ------------------------------- Evaluation: CUDA_VISIBLE_DEVICES=0,1 python finetuning.py --evaluate --eval_type fast --test_eps 1 --resume XXXX/model_best.pth.tar CUDA_VISIBLE_DEVICES=0,1 python finetuning.py --evaluate --eval_type fast --test_eps 1 --resume Source_PT/TeCoAmodel_best.pth.tar """ def parse_option(): parser = argparse.ArgumentParser('Visual Prompting for CLIP') parser.add_argument('--print_freq', type=int, default=2000, help='print frequency') parser.add_argument('--save_freq', type=int, default=50, help='save frequency') parser.add_argument('--validate_freq', type=int, default=1, help='validate frequency') parser.add_argument('--batch_size', type=int, default=256, help='batch_size') parser.add_argument('--num_workers', type=int, default=32, help='num of workers to use') parser.add_argument('--epochs', type=int, default=10, help='number of training epoch5s') parser.add_argument("--mix_alpha", type=float, default=-1, help="interpolation") # optimization parser.add_argument('--optim', type=str, default='sgd', help='optimizer to use') parser.add_argument('--learning_rate', type=float, default=1e-5, ## Change from 1e-7 to 1e-5 help='learning rate') parser.add_argument("--weight_decay", type=float, default=0, help="weight decay") parser.add_argument("--warmup", type=int, default=1000, help="number of steps to warmup for") parser.add_argument('--momentum', type=float, default=0.9, help='momentum') parser.add_argument('--train_eps', type=float, default=2, help='momentum') parser.add_argument('--train_numsteps', type=int, default=5) parser.add_argument('--train_stepsize', type=int, default=1) parser.add_argument('--test_eps', type=float, default=2, help='momentum') parser.add_argument('--test_numsteps', type=int, default=20) parser.add_argument('--test_stepsize', type=int, default=1) parser.add_argument('--patience', type=int, default=1000) # model parser.add_argument('--model', type=str, default='clip') parser.add_argument('--imagenet_root', type=str, default='temp') parser.add_argument('--arch', type=str, default='vit_b32') parser.add_argument('--method', type=str, default='null_patch', choices=['null_patch'], help='choose visual prompting method') parser.add_argument('--name', type=str, default='') parser.add_argument('--prompt_size', type=int, default=30, help='size for visual prompts') parser.add_argument('--add_prompt_size', type=int, default=0, help='size for additional visual prompts') # dataset parser.add_argument('--root', type=str, default='/home/data1/junhao/datasets/', help='dataset') parser.add_argument('--dataset', type=str, default='ImageNet', help='Pre-training Dataset: cifar10|cifar100|ImageNet') parser.add_argument('--image_size', type=int, default=224, help='image size') # other parser.add_argument('--seed', type=int, default=None, help='seed for initializing training') parser.add_argument('--model_dir', type=str, default='../save_ckpts', help='path to save models') parser.add_argument('--image_dir', type=str, default='./save/images', help='path to save images') parser.add_argument('--filename', type=str, default=None, help='filename to save') parser.add_argument('--trial', type=int, default=1, help='number of trials') parser.add_argument('--gpu', type=int, default=None, help='gpu to use') parser.add_argument('--debug', action='store_true') parser.add_argument('--VPbaseline', action='store_true') parser.add_argument('--CW', action='store_true') parser.add_argument('--autoattack', action='store_true') parser.add_argument('--train_class_count', type=int, default=90) parser.add_argument('--last_num_ft', type=int, default=-1) parser.add_argument('--noimginprop', action='store_true') parser.add_argument('--exp_name', type=str, default=None) # Augmentation parser.add_argument('--aug_type', type=str, default='Vanilla', help='Vanilla|Vanilla_Flip|Resizecrop|Resizecrop_Flip|Autoaug') # Evaluation parser.add_argument('--resume', type=str, default=None, help='path to resume from checkpoint') parser.add_argument('--evaluate', default=False, action="store_true", help='evaluate model test set') parser.add_argument('--eval_type', type=str, default="fast", help='fast|full') # Extra modules parser.add_argument('--W_Pred_Align', type=float, default=0.0, help='Prediction alignment between clean and adv logits') parser.add_argument('--W_Nat_CE', type=float, default=0.0, help='Natural classification of clean logit') parser.add_argument('--W_Pred_Align_Ori', type=float, default=0.0, help='Prediction alignment between adv logits to the original clip-clean logits') args = parser.parse_args() if args.exp_name is not None: args.filename = args.exp_name else: args.filename = '{}_{}_{}_{}_{}_{}_{}_lr_{}_decay_{}_bsz_{}_warmup_{}_trial_{}_addp_{}'. \ format(args.name, args.method, args.prompt_size, args.dataset, args.model, args.arch, args.optim, args.learning_rate, args.weight_decay, args.batch_size, args.warmup, args.trial, args.add_prompt_size) return args best_acc1 = 0 device = "cuda" if torch.cuda.is_available() else "cpu" CIFAR100_MEAN = (0.48145466, 0.4578275, 0.40821073) CIFAR100_STD = (0.26862954, 0.26130258, 0.27577711) mu = torch.tensor(CIFAR100_MEAN).view(3, 1, 1).cuda() std = torch.tensor(CIFAR100_STD).view(3, 1, 1).cuda() def normalize(X): return (X - mu) / std def clip_img_preprocessing(X): img_size = 224 X = torch.nn.functional.upsample(X, size=(img_size, img_size), mode='bicubic') X = normalize(X) return X # for multiGPU clip def create_logits(x1, x2, logit_scale): x1 = x1 / x1.norm(dim=-1, keepdim=True) x2 = x2 / x2.norm(dim=-1, keepdim=True) # cosine similarity as logits logits_per_x1 = logit_scale * x1 @ x2.t() logits_per_x2 = logit_scale * x2 @ x1.t() # shape = [global_batch_size, global_batch_size] return logits_per_x1, logits_per_x2 class TextCLIP(nn.Module): def __init__(self, model): super(TextCLIP, self).__init__() self.model = model def forward(self, text): return self.model.encode_text(text) class ImageCLIP(nn.Module): def __init__(self, model): super(ImageCLIP, self).__init__() self.model = model def forward(self, image, prompt_token=None): return self.model.encode_image(image, prompt_token) ### # alpha_test = 1. / 255 # attack_iters_test = 5 # # epsilon = 2./255 upper_limit, lower_limit = 1, 0 def main(): global best_acc1, device args = parse_option() args.train_eps = args.train_eps / 255. args.test_eps = args.test_eps / 255. args.train_stepsize = args.train_stepsize / 255. args.test_stepsize = args.test_stepsize / 255. if args.resume is not None: args.resume = os.path.join("../save_ckpts", args.resume) print(args) if args.seed is not None: random.seed(args.seed) torch.manual_seed(args.seed) cudnn.deterministic = True else: cudnn.benchmark = True import socket if socket.gethostname() == 'junhao': args.root = '/home/data1/junhao/datasets/' elif socket.gethostname() == 'ai-planning-p4de-02': args.root = '/data_3/teddy_research/datasets_jh/' # if args.imagenet_root is not None: imagenet_root = os.path.join(args.root, "ImageNet") imgnet_full = imagenet_root # create model # add_prompt_len = args.add_prompt_size add_prompt_len = 0 model, preprocess = clip.load('ViT-B/32', device, jit=False, prompt_len=add_prompt_len) model_text, model_image = None, None convert_models_to_fp32(model) model = torch.nn.DataParallel(model) # .to(device) model.eval() original_model = None if args.W_Pred_Align_Ori > 0.0: original_model, preprocess = clip.load('ViT-B/32', device, jit=False, prompt_len=add_prompt_len) convert_models_to_fp32(original_model) original_model = torch.nn.DataParallel(original_model) # .to(device) original_model.eval() prompter = NullPrompter() # .to(device) add_prompter = TokenPrompter(add_prompt_len) # .to(device) prompter = torch.nn.DataParallel(prompter).cuda() add_prompter = torch.nn.DataParallel(add_prompter).cuda() # define criterion and optimizer # we finetune the image module parameters only if args.last_num_ft == -1: optimizer = torch.optim.SGD(model.module.visual.parameters(), lr=args.learning_rate, momentum=args.momentum, weight_decay=args.weight_decay) else: optimizer = torch.optim.SGD(list(model.module.visual.parameters())[-args.last_num_ft:], lr=args.learning_rate, momentum=args.momentum, weight_decay=args.weight_decay) criterion = torch.nn.CrossEntropyLoss().to(device) args.start_epoch = 0 # optionally resume from a checkpoint if args.resume: if os.path.isfile(args.resume): print("=> loading checkpoint '{}'".format(args.resume)) if args.gpu is None: checkpoint = torch.load(args.resume) else: # Map model to be loaded to specified single gpu. loc = 'cuda:{}'.format(args.gpu) checkpoint = torch.load(args.resume, map_location=loc) args.start_epoch = checkpoint['epoch'] best_acc1 = checkpoint['best_acc1'] if args.gpu is not None: # best_acc1 may be from a checkpoint from a different GPU best_acc1 = best_acc1.to(args.gpu) if args.mix_alpha > 0: alpha = args.mix_alpha # model1, preprocess = clip.load('ViT-B/32', device, jit=False, prompt_len=add_prompt_len) # model2, preprocess = clip.load('ViT-B/32', device, jit=False, prompt_len=add_prompt_len) # model1 = torch.nn.DataParallel(model1) # model2 = torch.nn.DataParallel(model2) checkpoint_ori = torch.load('original_clip.pth.tar') theta_ori = checkpoint_ori['vision_encoder_state_dict'] theta_rob = checkpoint['vision_encoder_state_dict'] theta = { key: (1 - alpha) * theta_ori[key] + alpha * theta_rob[key] for key in theta_ori.keys() } model.module.visual.load_state_dict(theta) else: model.module.visual.load_state_dict(checkpoint['vision_encoder_state_dict']) optimizer.load_state_dict(checkpoint['optimizer']) # prompter.load_state_dict(checkpoint['state_dict']) # add_prompter.load_state_dict(checkpoint['add_prompter']) print("=> loaded checkpoint '{}' (epoch {})" .format(args.resume, checkpoint['epoch'])) else: print("=> no checkpoint found at '{}'".format(args.resume)) # create data template = 'This is a photo of a {}' print(f'template: {template}') # # print('model.module.visual', list(model.module.visual.parameters())[-3:]) # for each in list(model.module.visual.parameters()): # print(each.size()) # exit() # print(preprocess, 'preprocess') # exit() # TODO: we can train on cifar10 and test on cifar10, 100 in zero shot way, to see if generalize. preprocess = transforms.Compose([ # transforms.RandomHorizontalFlip(), # transforms.RandomRotation(15), # TODO: may use later transforms.ToTensor() ]) preprocess224 = transforms.Compose([ transforms.Resize(256), transforms.CenterCrop(224), # transforms.RandomHorizontalFlip(), # transforms.RandomRotation(15), # TODO: may use later transforms.ToTensor() ]) preprocess224_interpolate = transforms.Compose([ transforms.Resize((224, 224)), # transforms.RandomHorizontalFlip(), # transforms.RandomRotation(15), # TODO: may use later transforms.ToTensor() ]) ############################ Augmentation ############################ preprocess224_vanilla_flip = transforms.Compose([ transforms.Resize(256), transforms.CenterCrop(224), transforms.RandomHorizontalFlip(), transforms.ToTensor() ]) preprocess224_resizecrop = transforms.Compose([ transforms.RandomResizedCrop(224), transforms.ToTensor() ]) preprocess224_resizecrop_flip = transforms.Compose([ transforms.RandomResizedCrop(224), transforms.RandomHorizontalFlip(), transforms.ToTensor() ]) preprocess_autoaug = transforms.Compose([ transforms.RandomResizedCrop(224), transforms.RandomHorizontalFlip(), ImageNetPolicy(), transforms.ToTensor() ]) # Vanilla|Vanilla_Flip|Resizecrop|Resizecrop_Flip|Autoaug if args.aug_type == 'Vanilla': IN_aug_type = preprocess224 elif args.aug_type == 'Vanilla_Flip': IN_aug_type = preprocess224_vanilla_flip elif args.aug_type == 'Resizecrop': IN_aug_type = preprocess224_resizecrop elif args.aug_type == 'Resizecrop_Flip': IN_aug_type = preprocess224_resizecrop_flip elif args.aug_type == 'Autoaug': IN_aug_type = preprocess_autoaug ############################ Augmentation ############################ if args.dataset == 'cifar100': print('hi') train_dataset = CIFAR100(args.root, transform=preprocess, download=True, train=True) val_dataset = CIFAR100(args.root, transform=preprocess, download=True, train=False) elif args.dataset == 'cifar10': train_dataset = CIFAR10(args.root, transform=preprocess, download=True, train=True) val_dataset = CIFAR10(args.root, transform=preprocess, download=True, train=False) elif args.dataset == 'ImageNet': train_dataset = torchvision.datasets.ImageFolder( os.path.join(imagenet_root, 'train'), transform=IN_aug_type ) val_dataset_list = [] val_dataset_name = ['StanfordCars', 'Food101', 'PCAM', 'cifar100', 'oxfordpet', 'flowers102', 'Country211', 'dtd', 'EuroSAT', 'fgvc_aircraft', 'ImageNet', 'cifar10', 'SUN397'] if args.evaluate: if args.eval_type == 'fast': val_dataset_name = ['ImageNet', 'SUN397', 'Food101', 'flowers102', 'Caltech101', 'Caltech256'] elif args.eval_type == 'full': val_dataset_name = ['ImageNet', 'cifar10', 'STL10', 'cifar100', 'SUN397', 'StanfordCars', 'Food101', 'oxfordpet', 'flowers102', 'dtd', 'EuroSAT', 'fgvc_aircraft', 'PCAM', 'Caltech101', 'Caltech256'] else: val_dataset_name = ['cifar10', 'cifar100', 'dtd', 'EuroSAT'] for each in val_dataset_name: if each == 'cifar10': val_dataset_list.append(CIFAR10(args.root, transform=preprocess, download=True, train=False)) elif each == 'cifar100': val_dataset_list.append(CIFAR100(args.root, transform=preprocess, download=True, train=False)) elif each == 'Caltech101': val_dataset_list.append(Caltech101(args.root, target_type='category', transform=preprocess224, download=True)) elif each == 'PCAM': val_dataset_list.append(PCAM(args.root, split='test', transform=preprocess224, download=True)) elif each == 'STL10': val_dataset_list.append(STL10(args.root, split='test', transform=preprocess, download=True)) elif each == 'SUN397': val_dataset_list.append(SUN397(args.root, transform=preprocess224, download=True)) elif each == 'StanfordCars': val_dataset_list.append(StanfordCars(args.root, split='test', transform=preprocess224, download=True)) elif each == 'Food101': val_dataset_list.append(Food101(args.root, split='test', transform=preprocess224, download=True)) elif each == 'oxfordpet': val_dataset_list.append(OxfordIIITPet(args.root, split='test', transform=preprocess224, download=True)) elif each == 'EuroSAT': val_dataset_list.append(EuroSAT(args.root, transform=preprocess224, download=True)) elif each == 'Caltech256': val_dataset_list.append(Caltech256(args.root, transform=preprocess224, download=True)) # elif each == 'FER2013': # val_dataset_list.append(OxfordIIITPet(args.root, split='test', # transform=preprocess224, download=True)) elif each == 'flowers102': val_dataset_list.append(Flowers102(args.root, split='test', transform=preprocess224, download=True)) elif each == 'Country211': val_dataset_list.append(Country211(args.root, split='test', transform=preprocess224, download=True)) elif each == 'dtd': val_dataset_list.append(DTD(args.root, split='test', transform=preprocess224, download=True)) elif each == 'fgvc_aircraft': val_dataset_list.append(FGVCAircraft(args.root, split='test', transform=preprocess224, download=True)) elif each == 'ImageNet': val_dataset_list.append(torchvision.datasets.ImageFolder( os.path.join(imgnet_full, 'val'), transform=preprocess224)) # val_dataset_list.append(torchvision.datasets.ImageNet( # root=imagenet_root, # split='val', # transform=preprocess224)) train_sampler = None val_sampler = None train_loader = DataLoader(train_dataset, batch_size=args.batch_size, pin_memory=True, num_workers=args.num_workers, shuffle=True, sampler=train_sampler) val_loader_list = [DataLoader(each, batch_size=args.batch_size, pin_memory=True, num_workers=args.num_workers, shuffle=False, sampler=val_sampler) for each in val_dataset_list] class_names = train_dataset.classes if args.dataset == 'ImageNet': from utils import load_imagenet_folder2name folder2name = load_imagenet_folder2name('imagenet_classes_names.txt') new_class_names = [] for each in class_names: new_class_names.append(folder2name[each]) class_names = new_class_names class_names = refine_classname(class_names) texts_train = [template.format(label) for label in class_names] texts_list = [] for cnt, each in enumerate(val_dataset_list): if hasattr(each, 'clip_prompts'): texts_tmp = each.clip_prompts else: class_names = each.classes if val_dataset_name[cnt] == 'ImageNet': from utils import load_imagenet_folder2name folder2name = load_imagenet_folder2name('imagenet_classes_names.txt') new_class_names = [] for class_name in class_names: new_class_names.append(folder2name[class_name]) class_names = new_class_names class_names = refine_classname(class_names) texts_tmp = [template.format(label) for label in class_names] texts_list.append(texts_tmp) assert len(texts_list) == len(val_dataset_list) scaler = GradScaler() total_steps = len(train_loader) * args.epochs scheduler = cosine_lr(optimizer, args.learning_rate, args.warmup, total_steps) # make dir refined_template = template.lower().replace(' ', '_') # args.filename = f'{args.filename}_template_{refined_template}' args.model_folder = os.path.join(args.model_dir, args.filename) if not os.path.isdir(args.model_folder): os.makedirs(args.model_folder) # wandb # if args.use_wandb: # wandb.init(project='Visual Prompting') # wandb.config.update(args) # wandb.run.name = args.filename # wandb.watch(prompter, criterion, log='all', log_freq=10) if args.evaluate: acc1_mean = validate(val_loader_list, val_dataset_name, texts_list, model, model_text, model_image, prompter, add_prompter, criterion, args) return epochs_since_improvement = 0 for epoch in range(args.start_epoch, args.epochs): # train for one epoch train(train_loader, texts_train, model, original_model, model_text, model_image, prompter, add_prompter, optimizer, scheduler, criterion, scaler, epoch, args) # evaluate on validation set if epoch % args.validate_freq == 0: acc1_mean = validate(val_loader_list, val_dataset_name, texts_list, model, model_text, model_image, prompter, add_prompter, criterion, args) # remember best acc@1 and save checkpoint is_best = acc1_mean > best_acc1 best_acc1 = max(acc1_mean, best_acc1) save_checkpoint({ 'epoch': epoch + 1, 'state_dict': prompter.state_dict(), 'add_prompter': add_prompter.state_dict(), 'vision_encoder_state_dict': model.module.visual.state_dict(), 'best_acc1': best_acc1, 'optimizer': optimizer.state_dict(), }, args, is_best=is_best) if is_best: epochs_since_improvement = 0 else: epochs_since_improvement += 1 print(f"There's no improvement for {epochs_since_improvement} epochs.") if epochs_since_improvement >= args.patience: print("The training halted by early stopping criterion.") break # wandb.run.finish() def clamp(X, lower_limit, upper_limit): return torch.max(torch.min(X, upper_limit), lower_limit) from utils import one_hot_embedding def attack_CW(prompter, model, model_text, model_image, add_prompter, criterion, X, target, text_tokens, alpha, attack_iters, norm, restarts=1, early_stop=True, epsilon=0): delta = torch.zeros_like(X).cuda() if norm == "l_inf": delta.uniform_(-epsilon, epsilon) elif norm == "l_2": delta.normal_() d_flat = delta.view(delta.size(0), -1) n = d_flat.norm(p=2, dim=1).view(delta.size(0), 1, 1, 1) r = torch.zeros_like(n).uniform_(0, 1) delta *= r / n * epsilon else: raise ValueError delta = clamp(delta, lower_limit - X, upper_limit - X) delta.requires_grad = True for _ in range(attack_iters): # output = model(normalize(X )) prompted_images = prompter(clip_img_preprocessing(X + delta)) prompt_token = add_prompter() output, _ = multiGPU_CLIP(model_image, model_text, model, prompted_images, text_tokens, prompt_token) num_class = output.size(1) label_mask = one_hot_embedding(target, num_class) label_mask = label_mask.cuda() correct_logit = torch.sum(label_mask * output, dim=1) wrong_logit, _ = torch.max((1 - label_mask) * output - 1e4 * label_mask, axis=1) # loss = criterion(output, target) loss = - torch.sum(F.relu(correct_logit - wrong_logit + 50)) loss.backward() grad = delta.grad.detach() d = delta[:, :, :, :] g = grad[:, :, :, :] x = X[:, :, :, :] if norm == "l_inf": d = torch.clamp(d + alpha * torch.sign(g), min=-epsilon, max=epsilon) elif norm == "l_2": g_norm = torch.norm(g.view(g.shape[0], -1), dim=1).view(-1, 1, 1, 1) scaled_g = g / (g_norm + 1e-10) d = (d + scaled_g * alpha).view(d.size(0), -1).renorm(p=2, dim=0, maxnorm=epsilon).view_as(d) d = clamp(d, lower_limit - x, upper_limit - x) delta.data[:, :, :, :] = d delta.grad.zero_() return delta def attack_CW_noprompt(prompter, model, model_text, model_image, criterion, X, target, text_tokens, alpha, attack_iters, norm, restarts=1, early_stop=True, epsilon=0): delta = torch.zeros_like(X).cuda() if norm == "l_inf": delta.uniform_(-epsilon, epsilon) elif norm == "l_2": delta.normal_() d_flat = delta.view(delta.size(0), -1) n = d_flat.norm(p=2, dim=1).view(delta.size(0), 1, 1, 1) r = torch.zeros_like(n).uniform_(0, 1) delta *= r / n * epsilon else: raise ValueError delta = clamp(delta, lower_limit - X, upper_limit - X) delta.requires_grad = True for _ in range(attack_iters): # output = model(normalize(X )) _images = clip_img_preprocessing(X + delta) # output, _ = model(_images, text_tokens) output, _ = multiGPU_CLIP(model_image, model_text, model, _images, text_tokens, None) num_class = output.size(1) label_mask = one_hot_embedding(target, num_class) label_mask = label_mask.cuda() correct_logit = torch.sum(label_mask * output, dim=1) wrong_logit, _ = torch.max((1 - label_mask) * output - 1e4 * label_mask, axis=1) # loss = criterion(output, target) loss = - torch.sum(F.relu(correct_logit - wrong_logit + 50)) loss.backward() grad = delta.grad.detach() d = delta[:, :, :, :] g = grad[:, :, :, :] x = X[:, :, :, :] if norm == "l_inf": d = torch.clamp(d + alpha * torch.sign(g), min=-epsilon, max=epsilon) elif norm == "l_2": g_norm = torch.norm(g.view(g.shape[0], -1), dim=1).view(-1, 1, 1, 1) scaled_g = g / (g_norm + 1e-10) d = (d + scaled_g * alpha).view(d.size(0), -1).renorm(p=2, dim=0, maxnorm=epsilon).view_as(d) d = clamp(d, lower_limit - x, upper_limit - x) delta.data[:, :, :, :] = d delta.grad.zero_() return delta def attack_pgd(prompter, model, model_text, model_image, add_prompter, criterion, X, target, text_tokens, alpha, attack_iters, norm, restarts=1, early_stop=True, epsilon=0): delta = torch.zeros_like(X).cuda() if norm == "l_inf": delta.uniform_(-epsilon, epsilon) elif norm == "l_2": delta.normal_() d_flat = delta.view(delta.size(0), -1) n = d_flat.norm(p=2, dim=1).view(delta.size(0), 1, 1, 1) r = torch.zeros_like(n).uniform_(0, 1) delta *= r / n * epsilon else: raise ValueError delta = clamp(delta, lower_limit - X, upper_limit - X) delta.requires_grad = True for _ in range(attack_iters): # output = model(normalize(X )) prompted_images = prompter(clip_img_preprocessing(X + delta)) prompt_token = add_prompter() output, _ = multiGPU_CLIP(model_image, model_text, model, prompted_images, text_tokens, prompt_token) loss = criterion(output, target) loss.backward() grad = delta.grad.detach() d = delta[:, :, :, :] g = grad[:, :, :, :] x = X[:, :, :, :] if norm == "l_inf": d = torch.clamp(d + alpha * torch.sign(g), min=-epsilon, max=epsilon) elif norm == "l_2": g_norm = torch.norm(g.view(g.shape[0], -1), dim=1).view(-1, 1, 1, 1) scaled_g = g / (g_norm + 1e-10) d = (d + scaled_g * alpha).view(d.size(0), -1).renorm(p=2, dim=0, maxnorm=epsilon).view_as(d) d = clamp(d, lower_limit - x, upper_limit - x) delta.data[:, :, :, :] = d delta.grad.zero_() return delta def attack_pgd_noprompt(prompter, model, model_text, model_image, criterion, X, target, text_tokens, alpha, attack_iters, norm, restarts=1, early_stop=True, epsilon=0): delta = torch.zeros_like(X).cuda() if norm == "l_inf": delta.uniform_(-epsilon, epsilon) elif norm == "l_2": delta.normal_() d_flat = delta.view(delta.size(0), -1) n = d_flat.norm(p=2, dim=1).view(delta.size(0), 1, 1, 1) r = torch.zeros_like(n).uniform_(0, 1) delta *= r / n * epsilon else: raise ValueError delta = clamp(delta, lower_limit - X, upper_limit - X) delta.requires_grad = True for _ in range(attack_iters): _images = clip_img_preprocessing(X + delta) output, _ = multiGPU_CLIP(model_image, model_text, model, _images, text_tokens, None) loss = criterion(output, target) loss.backward() grad = delta.grad.detach() d = delta[:, :, :, :] g = grad[:, :, :, :] x = X[:, :, :, :] if norm == "l_inf": d = torch.clamp(d + alpha * torch.sign(g), min=-epsilon, max=epsilon) elif norm == "l_2": g_norm = torch.norm(g.view(g.shape[0], -1), dim=1).view(-1, 1, 1, 1) scaled_g = g / (g_norm + 1e-10) d = (d + scaled_g * alpha).view(d.size(0), -1).renorm(p=2, dim=0, maxnorm=epsilon).view_as(d) d = clamp(d, lower_limit - x, upper_limit - x) delta.data[:, :, :, :] = d delta.grad.zero_() return delta def attack_auto(model, images, target, text_tokens, prompter, add_prompter, attacks_to_run=['apgd-ce', 'apgd-dlr'], epsilon=0): forward_pass = functools.partial( multiGPU_CLIP_image_logits, model=model, text_tokens=text_tokens, prompter=None, add_prompter=None ) adversary = AutoAttack(forward_pass, norm='Linf', eps=epsilon, version='standard', verbose=False) adversary.attacks_to_run = attacks_to_run x_adv = adversary.run_standard_evaluation(images, target, bs=images.shape[0]) return x_adv def multiGPU_CLIP_image_logits(images, model, text_tokens, prompter=None, add_prompter=None): image_tokens = clip_img_preprocessing(images) prompt_token = None if add_prompter is None else add_prompter() if prompter is not None: image_tokens = prompter(image_tokens) return multiGPU_CLIP(None, None, model, image_tokens, text_tokens, prompt_token=prompt_token)[0] def multiGPU_CLIP(model_image, model_text, model, images, text_tokens, prompt_token=None): if prompt_token is not None: bs = images.size(0) prompt_token = prompt_token.repeat(bs, 1, 1) img_embed, scale_text_embed = model(images, text_tokens, prompt_token) logits_per_image = img_embed @ scale_text_embed.t() logits_per_text = scale_text_embed @ img_embed.t() return logits_per_image, logits_per_text def train(train_loader, texts, model, original_model, model_text, model_image, prompter, add_prompter, optimizer, scheduler, criterion, scaler, epoch, args): batch_time = AverageMeter('Time', ':6.3f') data_time = AverageMeter('Data', ':6.3f') losses = AverageMeter('Loss', ':.4e') top1 = AverageMeter('Acc@1', ':6.2f') progress = ProgressMeter( len(train_loader), [batch_time, data_time, losses, top1], prefix="Epoch: [{}]".format(epoch)) # switch to train mode model.module.visual.train() num_batches_per_epoch = len(train_loader) alpha = args.train_stepsize attack_iters = args.train_numsteps # print('text token', texts) end = time.time() for i, (images, target) in enumerate(tqdm(train_loader, ncols = 80)): # measure data loading time data_time.update(time.time() - end) BATCH_SIZE = images.size(0) # print('bs', BATCH_SIZE) # adjust learning rate step = num_batches_per_epoch * epoch + i scheduler(step) optimizer.zero_grad() images = images.to(device) target = target.to(device) text_tokens = clip.tokenize(texts).to(device) # print(images.min(), images.max()) # with automatic mixed precision with autocast(): loss_Pred_Align = 0.0 loss_Nat_CE = 0.0 loss_Pred_Align_Ori = 0.0 output_Inat_Tnat = None if not args.VPbaseline: delta = attack_pgd(prompter, model, model_text, model_image, add_prompter, criterion, images, target, text_tokens, alpha, attack_iters, 'l_inf', epsilon=args.train_eps) # print('delta', delta.min(), delta.max()) tmp = clip_img_preprocessing(images + delta) else: tmp = clip_img_preprocessing(images) prompted_images = prompter(tmp) prompt_token = None # Compute logits_image(256, 1000), logits_text(1000, 256) (Image-Text Alignment) output_Iadv_Tnat, _ = multiGPU_CLIP(model_image, model_text, model, prompted_images, text_tokens, prompt_token) if args.W_Pred_Align > 0.0: criterion_KL = nn.KLDivLoss(reduction='batchmean').to(device) tmp_nat = clip_img_preprocessing(images) prompted_nat_images = prompter(tmp_nat) output_Inat_Tnat, _ = multiGPU_CLIP(model_image, model_text, model, prompted_nat_images, text_tokens, prompt_token) loss_Pred_Align = criterion_KL(F.log_softmax(output_Iadv_Tnat, dim=1), F.softmax(output_Inat_Tnat, dim=1)) if args.W_Nat_CE > 0.0: if output_Inat_Tnat is None: tmp_nat = clip_img_preprocessing(images) prompted_nat_images = prompter(tmp_nat) output_Inat_Tnat, _ = multiGPU_CLIP(model_image, model_text, model, prompted_nat_images, text_tokens, prompt_token) loss_Nat_CE = criterion(output_Inat_Tnat, target) if args.W_Pred_Align_Ori > 0.0: criterion_KL = nn.KLDivLoss(reduction='batchmean').to(device) tmp_nat = clip_img_preprocessing(images) prompted_nat_images = prompter(tmp_nat) Ori_output_Inat_Tnat, _ = multiGPU_CLIP(model_image, model_text, original_model, prompted_nat_images, text_tokens, prompt_token) loss_Pred_Align_Ori = criterion_KL(F.log_softmax(output_Iadv_Tnat, dim=1), F.softmax(Ori_output_Inat_Tnat, dim=1)) loss = criterion(output_Iadv_Tnat, target) + args.W_Pred_Align * loss_Pred_Align + args.W_Nat_CE * loss_Nat_CE + args.W_Pred_Align_Ori * loss_Pred_Align_Ori scaler.scale(loss).backward() scaler.step(optimizer) scaler.update() # Note: we clamp to 4.6052 = ln(100), as in the original paper. model.module.logit_scale.data = torch.clamp(model.module.logit_scale.data, 0, 4.6052) # measure accuracy acc1 = accuracy(output_Iadv_Tnat, target, topk=(1,)) losses.update(loss.item(), images.size(0)) top1.update(acc1[0].item(), images.size(0)) # measure elapsed time batch_time.update(time.time() - end) end = time.time() if i % args.print_freq == 0 and i != 0: progress.display(i) if args.debug: break # break # if args.use_wandb: # wandb.log({ # 'training_loss': losses.avg, # 'training_acc': top1.avg # }) if i % args.save_freq == 0: save_checkpoint({ 'epoch': epoch + 1, 'state_dict': prompter.state_dict(), 'add_prompter': add_prompter.state_dict(), 'vision_encoder_state_dict': model.module.visual.state_dict(), 'best_acc1': best_acc1, 'optimizer': optimizer.state_dict(), }, args) return losses.avg, top1.avg # def validate(val_loader, texts, model, prompter, add_prompter, criterion, args): def validate(val_loader_list, val_dataset_name, texts_list, model, model_text, model_image, prompter, add_prompter, criterion, args): dataset_num = len(val_loader_list) acc_all_nat = [] acc_all_adv = [] test_stepsize = args.test_stepsize for cnt in range(dataset_num): val_loader = val_loader_list[cnt] texts = texts_list[cnt] dataset_name = val_dataset_name[cnt] binary = ['PCAM'] attacks_to_run=['apgd-ce', 'apgd-dlr'] if dataset_name in binary: attacks_to_run=['apgd-ce'] batch_time = AverageMeter('Time', ':6.3f') losses = AverageMeter('Loss', ':.4e') top1_org = AverageMeter('Original Acc@1', ':6.2f') top1_prompt = AverageMeter('Prompt Acc@1', ':6.2f') top1_adv_org = AverageMeter('Adv Original Acc@1', ':6.2f') top1_adv_prompt = AverageMeter('Adv Prompt Acc@1', ':6.2f') progress = ProgressMeter( len(val_loader), [batch_time, losses, top1_org, top1_adv_org], prefix=dataset_name + '_Validate: ') # switch to evaluation mode prompter.eval() add_prompter.eval() model.eval() # print(val_dataset_name, 'text token', texts_list) # end = time.time() for i, (images, target) in enumerate(tqdm(val_loader, ncols = 80)): if 'cifar' not in val_dataset_name: if i % 20 != 0 and not args.evaluate: continue images = images.to(device) target = target.to(device) text_tokens = clip.tokenize(texts).to(device) # print(images.size()) with autocast(): # clean images, with prompt and without prompt # compute output with torch.no_grad(): # prompt_token = add_prompter() prompt_token = None # output_prompt, _ = model(prompter(clip_img_preprocessing(images)), text_tokens, prompt_token) output_prompt, _ = multiGPU_CLIP(model_image, model_text, model, prompter(clip_img_preprocessing(images)), text_tokens, prompt_token) loss = criterion(output_prompt, target) # measure accuracy and record loss acc1 = accuracy(output_prompt, target, topk=(1,)) losses.update(loss.item(), images.size(0)) # top1_prompt.update(acc1[0].item(), images.size(0)) top1_org.update(acc1[0].item(), images.size(0)) torch.cuda.empty_cache() # generate adv example if args.CW: delta_prompt = attack_CW(prompter, model, model_text, model_image, add_prompter, criterion, images, target, text_tokens, test_stepsize, args.test_numsteps, 'l_inf', epsilon=args.test_eps) attacked_images = images + delta_prompt elif args.autoattack: attacked_images = attack_auto(model, images, target, text_tokens, None, None, epsilon=args.test_eps, attacks_to_run=attacks_to_run) else: delta_prompt = attack_pgd(prompter, model, model_text, model_image, add_prompter, criterion, images, target, text_tokens, test_stepsize, args.test_numsteps, 'l_inf', epsilon=args.test_eps) attacked_images = images + delta_prompt # compute output torch.cuda.empty_cache() with torch.no_grad(): prompt_token = add_prompter() # output_prompt_adv, _ = model(prompter(clip_img_preprocessing(images + delta_prompt)), text_tokens, prompt_token) output_prompt_adv, _ = multiGPU_CLIP(model_image, model_text, model, prompter(clip_img_preprocessing(attacked_images)), text_tokens, prompt_token) loss = criterion(output_prompt_adv, target) # bl attack torch.cuda.empty_cache() # measure accuracy and record loss acc1 = accuracy(output_prompt_adv, target, topk=(1,)) losses.update(loss.item(), images.size(0)) top1_adv_org.update(acc1[0].item(), images.size(0)) # top1_adv_prompt.update(acc1[0].item(), images.size(0)) # acc1 = accuracy(output_org_adv, target, topk=(1,)) # top1_adv_org.update(acc1[0].item(), images.size(0)) # measure elapsed time batch_time.update(time.time() - end) end = time.time() if i % args.print_freq == 0 and i != 0: progress.display(i) if args.debug: break torch.cuda.empty_cache() # print(dataset_name + ' * Adv Prompt Acc@1 {top1_adv_prompt.avg:.3f} Adv Original Acc@1 {top1_adv_org.avg:.3f} ' # '* Prompt Acc@1 {top1_prompt.avg:.3f} Original Acc@1 {top1_org.avg:.3f}' # .format(top1_adv_prompt=top1_adv_prompt, top1_adv_org=top1_adv_org, # top1_prompt=top1_prompt, top1_org=top1_org)) print(dataset_name + '--- Clean Acc.: {top1_org.avg:.2f} Adv Acc.: {top1_adv_org.avg:.2f}.' .format(top1_org=top1_org, top1_adv_org=top1_adv_org)) acc_all_nat.append(top1_org.avg) acc_all_adv.append(top1_adv_org.avg) # if args.use_wandb: # wandb.log({ # 'val_loss': losses.avg, # 'val_acc_prompt': top1_prompt.avg, # 'val_acc_org': top1_org.avg, # }) print('Average on all datasets --- Clean Acc.: {:.2f} Adv Acc.: {:.2f}.' .format(np.mean(acc_all_nat), np.mean(acc_all_adv))) return np.mean(acc_all_adv) if __name__ == '__main__': main()