import copy from torch.nn.modules import loss # from model.hash_model import DCMHT as DCMHT import os from tqdm import tqdm import torch import torch.nn as nn from torch.utils.data import DataLoader import torch.utils.data as data import scipy.io as scio import numpy as np from .base import TrainBase from torch.nn import functional as F from utils import get_args, calc_neighbor, cosine_similarity, euclidean_similarity, find_indices from utils.calc_utils import cal_map, cal_pr from dataset.dataloader import dataloader import clip import re from transformers import BertForMaskedLM from model.bert_tokenizer import BertTokenizer from transformers import AutoTokenizer device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu") filter_words = ['a', 'about', 'above', 'across', 'after', 'afterwards', 'again', 'against', 'ain', 'all', 'almost', 'alone', 'along', 'already', 'also', 'although', 'am', 'among', 'amongst', 'an', 'and', 'another', 'any', 'anyhow', 'anyone', 'anything', 'anyway', 'anywhere', 'are', 'aren', "aren't", 'around', 'as', 'at', 'back', 'been', 'before', 'beforehand', 'behind', 'being', 'below', 'beside', 'besides', 'between', 'beyond', 'both', 'but', 'by', 'can', 'cannot', 'could', 'couldn', "couldn't", 'd', 'didn', "didn't", 'doesn', "doesn't", 'don', "don't", 'down', 'due', 'during', 'either', 'else', 'elsewhere', 'empty', 'enough', 'even', 'ever', 'everyone', 'everything', 'everywhere', 'except', 'first', 'for', 'former', 'formerly', 'from', 'hadn', "hadn't", 'hasn', "hasn't", 'haven', "haven't", 'he', 'hence', 'her', 'here', 'hereafter', 'hereby', 'herein', 'hereupon', 'hers', 'herself', 'him', 'himself', 'his', 'how', 'however', 'hundred', 'i', 'if', 'in', 'indeed', 'into', 'is', 'isn', "isn't", 'it', "it's", 'its', 'itself', 'just', 'latter', 'latterly', 'least', 'll', 'may', 'me', 'meanwhile', 'mightn', "mightn't", 'mine', 'more', 'moreover', 'most', 'mostly', 'must', 'mustn', "mustn't", 'my', 'myself', 'namely', 'needn', "needn't", 'neither', 'never', 'nevertheless', 'next', 'no', 'nobody', 'none', 'noone', 'nor', 'not', 'nothing', 'now', 'nowhere', 'o', 'of', 'off', 'on', 'once', 'one', 'only', 'onto', 'or', 'other', 'others', 'otherwise', 'our', 'ours', 'ourselves', 'out', 'over', 'per', 'please', 's', 'same', 'shan', "shan't", 'she', "she's", "should've", 'shouldn', "shouldn't", 'somehow', 'something', 'sometime', 'somewhere', 'such', 't', 'than', 'that', "that'll", 'the', 'their', 'theirs', 'them', 'themselves', 'then', 'thence', 'there', 'thereafter', 'thereby', 'therefore', 'therein', 'thereupon', 'these', 'they', 'this', 'those', 'through', 'throughout', 'thru', 'thus', 'to', 'too', 'toward', 'towards', 'under', 'unless', 'until', 'up', 'upon', 'used', 've', 'was', 'wasn', "wasn't", 'we', 'were', 'weren', "weren't", 'what', 'whatever', 'when', 'whence', 'whenever', 'where', 'whereafter', 'whereas', 'whereby', 'wherein', 'whereupon', 'wherever', 'whether', 'which', 'while', 'whither', 'who', 'whoever', 'whole', 'whom', 'whose', 'why', 'with', 'within', 'without', 'won', "won't", 'would', 'wouldn', "wouldn't", 'y', 'yet', 'you', "you'd", "you'll", "you're", "you've", 'your', 'yours', 'yourself', 'yourselves', '.', '-', 'a the', '/', '?', 'some', '"', ',', 'b', '&', '!', '@', '%', '^', '*', '(', ')', "-", '-', '+', '=', '<', '>', '|', ':', ";", '～', '·'] filter_words = set(filter_words) def text_filter(text): text = re.findall(r"<|startoftext|>(.+)<|endoftext|>", text) text = text[2] text = re.sub(r'', ' ', text) return text class GoalFunctionStatus(object): SUCCEEDED = 0 # attack succeeded SEARCHING = 1 # In process of searching for a success FAILED = 2 # attack failed class GoalFunctionResult(object): goal_score = 1 def __init__(self, text, score=0, similarity=None): self.status = GoalFunctionStatus.SEARCHING self.text = text self.score = score self.similarity = similarity @property def score(self): return self.__score @score.setter def score(self, value): self.__score = value if value >= self.goal_score: self.status = GoalFunctionStatus.SUCCEEDED def __eq__(self, __o): return self.text == __o.text def __hash__(self): return hash(self.text) class Trainer(TrainBase): def __init__(self, rank=0): args = get_args() super(Trainer, self).__init__(args, rank) self.logger.info("dataset len: {}".format(len(self.train_loader.dataset))) image_mean, image_var = self.generate_mapping() self.image_mean = image_mean self.image_var = image_var self.device = rank self.clip_tokenizer = AutoTokenizer.from_pretrained("openai/clip-vit-base-patch32", model_max_length=77, truncation=True) self.bert_tokenizer = BertTokenizer.from_pretrained(self.args.text_encoder, do_lower_case=True) self.ref_net = BertForMaskedLM.from_pretrained(self.args.text_encoder).to(device) self.attack_thred = self.args.attack_thred # self.run() def _init_model(self): self.logger.info("init model.") model_clip, preprocess = clip.load(self.args.victim, device=device) self.model = model_clip self.model.eval() self.model.float() def _init_dataset(self): self.logger.info("init dataset.") self.logger.info(f"Using {self.args.dataset} dataset.") self.args.index_file = os.path.join("./dataset", self.args.dataset, self.args.index_file) self.args.caption_file = os.path.join("./dataset", self.args.dataset, self.args.caption_file) self.args.label_file = os.path.join("./dataset", self.args.dataset, self.args.label_file) train_data, query_data, retrieval_data = dataloader(captionFile=self.args.caption_file, indexFile=self.args.index_file, labelFile=self.args.label_file, maxWords=self.args.max_words, imageResolution=self.args.resolution, query_num=self.args.query_num, train_num=self.args.train_num, seed=self.args.seed) self.train_labels = train_data.get_all_label() self.query_labels = query_data.get_all_label() self.retrieval_labels = retrieval_data.get_all_label() self.args.retrieval_num = len(self.retrieval_labels) self.logger.info(f"query shape: {self.query_labels.shape}") self.logger.info(f"retrieval shape: {self.retrieval_labels.shape}") self.train_loader = DataLoader( dataset=train_data, batch_size=self.args.batch_size, num_workers=self.args.num_workers, pin_memory=True, shuffle=True ) self.query_loader = DataLoader( dataset=query_data, batch_size=self.args.batch_size, num_workers=self.args.num_workers, pin_memory=True, shuffle=True ) self.retrieval_loader = DataLoader( dataset=retrieval_data, batch_size=self.args.batch_size, num_workers=self.args.num_workers, pin_memory=True, shuffle=True ) self.train_data = train_data def _tokenize(self, text): words = text.split(' ') sub_words = [] keys = [] index = 0 for word in words: sub = self.bert_tokenizer.tokenize(word) sub_words += sub keys.append([index, index + len(sub)]) index += len(sub) return words, sub_words, keys def get_important_scores(self, text, origin_embeds, batch_size, max_length): # device = origin_embeds.device masked_words = self._get_masked(text) masked_texts = [' '.join(words) for words in masked_words] # list of text of masked words masked_embeds = [] for i in range(0, len(masked_texts), batch_size): masked_text_input = self.bert_tokenizer(masked_texts[i:i + batch_size], padding='max_length', truncation=True, max_length=max_length, return_tensors='pt').to( device) masked_embed = self.ref_net(masked_text_input.text_inputs, attention_mask=masked_text_input.attention_mask) masked_embeds.append(masked_embed) masked_embeds = torch.cat(masked_embeds, dim=0) criterion = torch.nn.KLDivLoss(reduction='none') import_scores = criterion(masked_embeds.log_softmax(dim=-1), origin_embeds.softmax(dim=-1).repeat(len(masked_texts), 1)) return import_scores.sum(dim=-1) def _get_masked(self, text): words = text.split(' ') len_text = len(words) masked_words = [] for i in range(len_text): masked_words.append(words[0:i] + ['[UNK]'] + words[i + 1:]) # list of words return masked_words def get_transformations(self, text, idx, substitutes): words = text.split(' ') trans_text = [] for sub in substitutes: words[idx] = sub trans_text.append(' '.join(words)) return trans_text def get_word_predictions(self, text): _, _, keys = self._tokenize(text) inputs = self.bert_tokenizer.encode_plus(text, add_special_tokens=True, max_length=self.args.max_words, truncation=True, return_tensors="pt") input_ids = inputs["input_ids"].to(self.device) attention_mask = inputs['attention_mask'] with torch.no_grad(): word_predictions = self.ref_net(input_ids)['logits'].squeeze(0) # (seq_len, vocab_size) # print(self.ref_net(input_ids)['logits'].shape) word_pred_scores_all, word_predictions = torch.topk(word_predictions, self.args.max_candidate, -1) word_predictions = word_predictions[1:-1, :] # remove [CLS] and [SEP] word_pred_scores_all = word_pred_scores_all[1:-1, :] return keys, word_predictions, word_pred_scores_all, attention_mask def get_bpe_substitutes(self, substitutes): # substitutes L, k substitutes = substitutes[0:12, 0:4] # maximum BPE candidates # find all possible candidates all_substitutes = [] for i in range(substitutes.size(0)): if len(all_substitutes) == 0: lev_i = substitutes[i] all_substitutes = [[int(c)] for c in lev_i] else: lev_i = [] for all_sub in all_substitutes: for j in substitutes[i]: lev_i.append(all_sub + [int(j)]) all_substitutes = lev_i # all substitutes: list of list of token-id (all candidates) cross_entropy_loss = nn.CrossEntropyLoss(reduction='none') word_list = [] all_substitutes = torch.tensor(all_substitutes) # [ N, L ] all_substitutes = all_substitutes[:24].to(self.device) N, L = all_substitutes.size() word_predictions = self.ref_net(all_substitutes)[0] # N L vocab-size ppl = cross_entropy_loss(word_predictions.view(N * L, -1), all_substitutes.view(-1)) # [ N*L ] ppl = torch.exp(torch.mean(ppl.view(N, L), dim=-1)) # N _, word_list = torch.sort(ppl) word_list = [all_substitutes[i] for i in word_list] final_words = [] for word in word_list: tokens = [self.bert_tokenizer.convert_ids_to_tokens(int(i)) for i in word] text = ' '.join([t.strip() for t in tokens]) final_words.append(text) return final_words def get_substitutes(self, substitutes, substitutes_score, threshold=3.0): ret = [] num_sub, _ = substitutes.size() if num_sub == 0: ret = [] elif num_sub == 1: for id, score in zip(substitutes[0], substitutes_score[0]): if threshold != 0 and score < threshold: break ret.append(self.bert_tokenizer.convert_ids_to_tokens(int(id))) elif self.args.enable_bpe: ret = self.get_bpe_substitutes(substitutes) return ret def filter_substitutes(self, substitues): ret = [] for word in substitues: if word.lower() in filter_words: continue if '##' in word: continue ret.append(word) return ret def get_goal_results(self, trans_texts, negetive_code, negetive_mean, negative_var, positive_code, positive_mean, positive_var, beta=10, temperature=0.05): # print(trans_texts) trans_feature = clip.tokenize(trans_texts,context_length=77,truncate=True).to(device) anchor = self.model.encode_text(trans_feature) batch_size=anchor.shape[0] loss1 = F.triplet_margin_with_distance_loss(anchor, positive_code.repeat(batch_size, 1), negetive_code.repeat(batch_size, 1), distance_function=nn.CosineSimilarity(), reduction='none') sim=F.cosine_similarity(anchor,positive_code.unsqueeze(0), dim=1, eps=1e-8).unsqueeze(1) negative_dist = (anchor - negetive_mean) ** 2 / (negative_var+ 1e-5) positive_dist = (anchor - positive_mean) ** 2 / (positive_var+ 1e-5) negatives = torch.exp(negative_dist / temperature) positives = torch.exp(positive_dist / temperature) loss = torch.log(positives / (positives + negatives)).mean(dim=1, keepdim=True) + beta * loss1 results = [] # print(loss.shape) for i in range(len(trans_texts)): if loss[i].shape[0] >1 or sim[i] >self.args.sim_threshold: continue results.append(GoalFunctionResult(trans_texts[i], score=loss[i], similarity=sim[i])) return results def generate_mapping(self): image_train = [] label_train = [] for image, text, label, index in self.train_loader: # raw_text=[self.clip_tokenizer.decode(token) for token in text] image = image.to(device, non_blocking=True) # print(self.model.vocab_size) temp_image = self.model.encode_image(image) image_train.append(temp_image.cpu().detach().numpy()) label_train.append(label.detach().numpy()) image_train = np.concatenate(image_train, axis=0) label_train = np.concatenate(label_train, axis=0) label_unipue = np.unique(label_train, axis=0) image_centroids = np.stack( [image_train[find_indices(label_train, label_unipue[i])].mean(axis=0) for i in range(len(label_unipue))], axis=0) image_var = np.stack( [image_train[find_indices(label_train, label_unipue[i])].var(axis=0) for i in range(len(label_unipue))], axis=0) image_representation = {} image_var_representation = {} for i, centroid in enumerate(label_unipue): image_representation[str(centroid.astype(int))] = image_centroids[i] image_var_representation[str(centroid.astype(int))] = image_var[i] return image_representation, image_var_representation def target_adv(self, raw_text, negetive_code, negetive_mean, negative_var, positive_code, positive_mean, positive_var, beta=10, temperature=0.05): # print(raw_text) keys, word_predictions, word_pred_scores_all, mask = self.get_word_predictions(raw_text) #clean state # clean_embeds=self.ref_net(bert_inputs.input_ids, attention_mask=bert_inputs.attention_mask) cur_result = GoalFunctionResult(raw_text) mask_idx = np.where(mask.cpu().numpy() == 1)[0] for idx in mask_idx: predictions = word_predictions[keys[idx][0]: keys[idx][1]] predictions_socre = word_pred_scores_all[keys[idx][0]: keys[idx][1]] substitutes = self.get_substitutes(predictions, predictions_socre) substitutes = self.filter_substitutes(substitutes) trans_texts = self.get_transformations(raw_text, idx, substitutes) if len(trans_texts) == 0: continue # loss function results = self.get_goal_results(trans_texts, negetive_code, negetive_mean, negative_var, positive_code, positive_mean, positive_var, beta, temperature) results = sorted(results, key=lambda x: x.score, reverse=True) if len(results) > 0 and results[0].score > cur_result.score: cur_result = results[0] else: continue if cur_result.status == GoalFunctionStatus.SUCCEEDED: max_similarity = cur_result.similarity if max_similarity is None: # similarity is not calculated continue for result in results[1:]: if result.status != GoalFunctionStatus.SUCCEEDED: break if result.similarity > max_similarity: max_similarity = result.similarity cur_result = result return cur_result if cur_result.status == GoalFunctionStatus.SEARCHING: cur_result.status = GoalFunctionStatus.FAILED return cur_result def train_epoch(self): # self.change_state(mode="valid") save_dir = os.path.join(self.args.save_dir, "adv_PR_cruve") all_loss = 0 times = 0 adv_codes = [] adv_label = [] for image, text, label, index in self.train_loader: self.global_step += 1 times += 1 print(times) image.float() image = image.to(self.rank, non_blocking=True) text = text.to(self.rank, non_blocking=True) negetive_mean = np.stack([self.image_mean[str(i.astype(int))] for i in label.detach().cpu().numpy()]) negative_var = np.stack([self.image_var[str(i.astype(int))] for i in label.detach().cpu().numpy()]) negetive_mean = torch.from_numpy(negetive_mean).to(self.rank, non_blocking=True) negative_var = torch.from_numpy(negative_var).to(self.rank, non_blocking=True) negetive_code = self.model.encode_image(image) #targeted sample np.random.seed(times) select_index = np.random.choice(len(self.train_data), size=self.args.batch_size) target_dataset = data.Subset(self.train_data, select_index) target_subset = torch.utils.data.DataLoader(target_dataset, batch_size=self.args.batch_size) target_image, _, target_label, _ = next(iter(target_subset)) target_image = target_image.to(self.rank, non_blocking=True) positive_mean = np.stack([self.image_mean[str(i.astype(int))] for i in target_label.detach().cpu().numpy()]) positive_var = np.stack([self.image_var[str(i.astype(int))] for i in target_label.detach().cpu().numpy()]) positive_mean = torch.from_numpy(positive_mean).to(self.rank, non_blocking=True) positive_var = torch.from_numpy(positive_var).to(self.rank, non_blocking=True) positive_code = self.model.encode_image(target_image) # print(self.clip_tokenizer.my_encode('This day is good!')) raw_text = [self.clip_tokenizer.convert_ids_to_tokens(token.cpu()) for token in text] raw_text = [text_filter(self.clip_tokenizer.convert_tokens_to_string(txt)) for txt in raw_text] final_texts=[] for i in range(self.args.batch_size): adv_txt=self.target_adv( raw_text[i], negetive_code[i], negetive_mean[i], negative_var[i], positive_code[i], positive_mean[i], positive_var[i]) final_texts.append(adv_txt.text) # final_adverse = self.target_adv( raw_text, negetive_code, negetive_mean, negative_var, # positive_code, positive_mean, positive_var) final_text = clip.tokenize(final_texts,context_length=77,truncate=True).to(self.rank, non_blocking=True) adv_code = self.model.encode_text(final_text) adv_codes.append(adv_code.cpu().detach().numpy()) adv_label.append(target_label.numpy()) adv_img = np.concatenate(adv_codes) adv_labels = np.concatenate(adv_label) _, retrieval_txt = self.get_code(self.retrieval_loader, self.args.retrieval_num) retrieval_txt = retrieval_txt.cpu().detach().numpy() retrieval_labels = self.retrieval_labels.numpy() mAP_t = cal_map(adv_img, adv_labels, retrieval_txt, retrieval_labels) self.logger.info(f">>>>>> MAP_t: {mAP_t}") result_dict = { 'adv_img': adv_img, 'r_txt': retrieval_txt, 'adv_l': adv_labels, 'r_l': retrieval_labels # 'q_l':query_labels # 'pr': pr, # 'pr_t': pr_t } scio.savemat(os.path.join(save_dir, str(self.args.victim).replace("/", "_") + "-adv-" + self.args.dataset + ".mat"), result_dict) self.logger.info(">>>>>> save all data!") def train(self): self.logger.info("Start train.") for epoch in range(self.args.epochs): self.train_epoch(epoch) self.valid(epoch) self.save_model(epoch) self.logger.info( f">>>>>>> FINISHED >>>>>> Best epoch, I-T: {self.best_epoch_i}, mAP: {self.max_mapi2t}, T-I: {self.best_epoch_t}, mAP: {self.max_mapt2i}") def make_hash_code(self, code: list) -> torch.Tensor: code = torch.stack(code) # print(code.shape) code = code.permute(1, 0, 2) hash_code = torch.argmax(code, dim=-1) hash_code[torch.where(hash_code == 0)] = -1 hash_code = hash_code.float() return hash_code def get_code(self, data_loader, length: int): img_buffer = torch.empty(length, self.args.output_dim, dtype=torch.float).to(self.rank) text_buffer = torch.empty(length, self.args.output_dim, dtype=torch.float).to(self.rank) for image, text, label, index in tqdm(data_loader): image = image.to(self.device, non_blocking=True) text = text.to(self.device, non_blocking=True) index = index.numpy() with torch.no_grad(): image_feature = self.model.encode_image(image) text_features = self.model.encode_text(text) img_buffer[index, :] = image_feature.detach() text_buffer[index, :] = text_features.detach() return img_buffer, text_buffer # img_buffer.to(self.rank), text_buffer.to(self.rank) def valid_attack(self, adv_images, texts, adv_labels): save_dir = os.path.join(self.args.save_dir, "adv_PR_cruve") os.makedirs(save_dir, exist_ok=True) def test(self, mode_name="i2t"): self.logger.info("Valid Clean.") save_dir = os.path.join(self.args.save_dir, "PR_cruve") os.makedirs(save_dir, exist_ok=True) query_img, query_txt = self.get_code(self.query_loader, self.args.query_num) retrieval_img, retrieval_txt = self.get_code(self.retrieval_loader, self.args.retrieval_num) query_img = query_img.cpu().detach().numpy() query_txt = query_txt.cpu().detach().numpy() retrieval_img = retrieval_img.cpu().detach().numpy() retrieval_txt = retrieval_txt.cpu().detach().numpy() query_labels = self.query_labels.numpy() retrieval_labels = self.retrieval_labels.numpy() mAPi2t = cal_map(query_img, query_labels, retrieval_txt, retrieval_labels) mAPt2i = cal_map(query_txt, query_labels, retrieval_img, retrieval_labels) # pr_i2t=cal_pr(retrieval_txt,query_img,query_labels,retrieval_labels) # pr_t2i=cal_pr(retrieval_img,query_txt,query_labels,retrieval_labels) self.max_mapt2i = max(self.max_mapt2i, mAPi2t) self.logger.info(f">>>>>> MAP(i->t): {mAPi2t}, MAP(t->i): {mAPt2i}") result_dict = { 'q_img': query_img, 'q_txt': query_txt, 'r_img': retrieval_img, 'r_txt': retrieval_txt, 'q_l': query_labels, 'r_l': retrieval_labels } scio.savemat(os.path.join(save_dir, str(self.args.victim).replace("/", "_") + "-ours-" + self.args.dataset + ".mat"), result_dict) self.logger.info(">>>>>> save all data!") def save_mat(self, query_img, query_txt, retrieval_img, retrieval_txt, mode_name="i2t"): save_dir = os.path.join(self.args.save_dir, "PR_cruve") os.makedirs(save_dir, exist_ok=True) query_img = query_img.cpu().detach().numpy() query_txt = query_txt.cpu().detach().numpy() retrieval_img = retrieval_img.cpu().detach().numpy() retrieval_txt = retrieval_txt.cpu().detach().numpy() query_labels = self.query_labels.numpy() retrieval_labels = self.retrieval_labels.numpy() result_dict = { 'q_img': query_img, 'q_txt': query_txt, 'r_img': retrieval_img, 'r_txt': retrieval_txt, 'q_l': query_labels, 'r_l': retrieval_labels } scio.savemat( os.path.join(save_dir, str(self.args.victim).replace("/", "_") + "-ours-" + self.args.dataset + "-" + mode_name + ".mat"), result_dict) self.logger.info(f">>>>>> save best {mode_name} data!")