添加 mlp_mix.py

.idea/.gitignore

@@ -0,0 +1,3 @@
+# Default ignored files
+/shelf/
+/workspace.xml

.idea/advclip.iml

@@ -0,0 +1,12 @@
+<?xml version="1.0" encoding="UTF-8"?>
+<module type="PYTHON_MODULE" version="4">
+  <component name="NewModuleRootManager">
+    <content url="file://$MODULE_DIR$" />
+    <orderEntry type="jdk" jdkName="torch2" jdkType="Python SDK" />
+    <orderEntry type="sourceFolder" forTests="false" />
+  </component>
+  <component name="PyDocumentationSettings">
+    <option name="format" value="PLAIN" />
+    <option name="myDocStringFormat" value="Plain" />
+  </component>
+</module>

.idea/inspectionProfiles/Project_Default.xml

@@ -0,0 +1,30 @@
+<component name="InspectionProjectProfileManager">
+  <profile version="1.0">
+    <option name="myName" value="Project Default" />
+    <inspection_tool class="PyPackageRequirementsInspection" enabled="true" level="WARNING" enabled_by_default="true">
+      <option name="ignoredPackages">
+        <value>
+          <list size="17">
+            <item index="0" class="java.lang.String" itemvalue="jax" />
+            <item index="1" class="java.lang.String" itemvalue="pyyaml" />
+            <item index="2" class="java.lang.String" itemvalue="flax" />
+            <item index="3" class="java.lang.String" itemvalue="tensorflow" />
+            <item index="4" class="java.lang.String" itemvalue="tensorboard" />
+            <item index="5" class="java.lang.String" itemvalue="jaxlib" />
+            <item index="6" class="java.lang.String" itemvalue="opencv-python" />
+            <item index="7" class="java.lang.String" itemvalue="Pillow" />
+            <item index="8" class="java.lang.String" itemvalue="transformers" />
+            <item index="9" class="java.lang.String" itemvalue="timm" />
+            <item index="10" class="java.lang.String" itemvalue="ruamel_yaml" />
+            <item index="11" class="java.lang.String" itemvalue="torch" />
+            <item index="12" class="java.lang.String" itemvalue="torchvision" />
+            <item index="13" class="java.lang.String" itemvalue="pandas" />
+            <item index="14" class="java.lang.String" itemvalue="scipy" />
+            <item index="15" class="java.lang.String" itemvalue="tqdm" />
+            <item index="16" class="java.lang.String" itemvalue="numpy" />
+          </list>
+        </value>
+      </option>
+    </inspection_tool>
+  </profile>
+</component>

.idea/inspectionProfiles/profiles_settings.xml

@@ -0,0 +1,6 @@
+<component name="InspectionProjectProfileManager">
+  <settings>
+    <option name="USE_PROJECT_PROFILE" value="false" />
+    <version value="1.0" />
+  </settings>
+</component>

.idea/misc.xml

@@ -0,0 +1,7 @@
+<?xml version="1.0" encoding="UTF-8"?>
+<project version="4">
+  <component name="Black">
+    <option name="sdkName" value="torch2" />
+  </component>
+  <component name="ProjectRootManager" version="2" project-jdk-name="torch2" project-jdk-type="Python SDK" />
+</project>

.idea/modules.xml

@@ -0,0 +1,8 @@
+<?xml version="1.0" encoding="UTF-8"?>
+<project version="4">
+  <component name="ProjectModuleManager">
+    <modules>
+      <module fileurl="file://$PROJECT_DIR$/.idea/advclip.iml" filepath="$PROJECT_DIR$/.idea/advclip.iml" />
+    </modules>
+  </component>
+</project>

.idea/vcs.xml

@@ -0,0 +1,6 @@
+<?xml version="1.0" encoding="UTF-8"?>
+<project version="4">
+  <component name="VcsDirectoryMappings">
+    <mapping directory="" vcs="Git" />
+  </component>
+</project>

auto_cifar10.py

@@ -0,0 +1,132 @@
+import warnings
+
+import torchvision.datasets
+
+warnings.filterwarnings('ignore')
+
+from PIL import Image
+import torch
+import timm
+import requests
+import numpy as np
+import torchvision.transforms as transforms
+from torch import nn
+from timm.data.constants import IMAGENET_DEFAULT_MEAN, IMAGENET_DEFAULT_STD
+from torch.utils.data import Dataset, DataLoader
+import copy
+
+from art.estimators.classification import PyTorchClassifier
+from art.data_generators import PyTorchDataGenerator
+from art.utils import load_cifar10
+from art.attacks.evasion import ProjectedGradientDescent ,AutoProjectedGradientDescent
+from art.defences.trainer import AdversarialTrainer
+
+model = timm.create_model("timm/vit_base_patch16_224.orig_in21k_ft_in1k", pretrained=False)
+model.head = nn.Linear(model.head.in_features, 10)
+state_dict = torch.load('/home/leewlving/.cache/torch/hub/checkpoints/vit_base_patch16_224_in21k_ft_cifar10.pth')
+model.load_state_dict(state_dict)
+# model.load_state_dict(
+#     torch.hub.load_state_dict_from_url(
+#         "https://huggingface.co/edadaltocg/vit_base_patch16_224_in21k_ft_cifar10/resolve/main/pytorch_model.bin",
+#         map_location="cuda",
+#         file_name="vit_base_patch16_224_in21k_ft_cifar10.pth",
+#     )
+# )
+model.eval()
+
+DEFAULT_MEAN = (0.485, 0.456, 0.406)
+DEFAULT_STD = (0.229, 0.224, 0.225)
+
+transform = transforms.Compose([
+    transforms.Resize(256, interpolation=3),
+    transforms.CenterCrop(224),
+    transforms.ToTensor(),
+    transforms.Normalize(DEFAULT_MEAN, DEFAULT_STD),
+])
+
+
+class CIFAR10_dataset(Dataset):
+    def __init__(self, data, targets, transform=None):
+        self.data = data
+        self.targets = torch.LongTensor(targets)
+        self.transform = transform
+
+    def __getitem__(self, index):
+        x = Image.fromarray(((self.data[index] * 255).round()).astype(np.uint8).transpose(1, 2, 0))
+        x = self.transform(x)
+        y = self.targets[index]
+        return x, y
+
+    def __len__(self):
+        return len(self.data)
+
+
+# (x_train, y_train), (x_test, y_test), min_pixel_value, max_pixel_value = load_cifar10()
+# print(max_pixel_value)
+# x_train = x_train.transpose(0, 3, 1, 2).astype("float32")
+# x_test = x_test.transpose(0, 3, 1, 2).astype("float32")
+train_dataset = torchvision.datasets.SVHN(root='./svhn',split='train',download=True,transform=transform)
+test_dataset= torchvision.datasets.SVHN(root='./svhn',split='test',download=True,transform=transform)
+# dataset = CIFAR10_dataset(x_train, y_train, transform=transform)
+dataloader = DataLoader(train_dataset, batch_size=64, shuffle=True)
+test_dataloader =DataLoader(test_dataset, batch_size=64, shuffle=False)
+
+opt = torch.optim.Adam(model.parameters(), lr=0.01)
+
+
+criterion = nn.CrossEntropyLoss()
+
+classifier = PyTorchClassifier(
+    model=model,
+    clip_values=(0.0, 1.0),
+    loss=criterion,
+    optimizer=opt,
+    input_shape=(3, 224, 224),
+    nb_classes=10,
+)
+
+
+attack= AutoProjectedGradientDescent(
+    classifier,
+    norm=np.inf,
+    eps=8.0 / 255.0,
+eps_step=2.0 / 255.0,
+    max_iter=10,
+    targeted=False,
+    batch_size=64,
+    verbose=False
+)
+# attack = ProjectedGradientDescent(
+#     classifier,
+#     norm=np.inf,
+#     eps=8.0 / 255.0,
+#     eps_step=2.0 / 255.0,
+#     max_iter=10,
+#     targeted=False,
+#     num_random_init=1,
+#     batch_size=64,
+#     verbose=False,
+# )
+
+trainer = AdversarialTrainer(
+    classifier, attack
+)
+art_datagen = PyTorchDataGenerator(iterator=dataloader, size=len(train_dataset), batch_size=64)
+
+trainer.fit_generator(art_datagen, nb_epochs=1)
+
+# for i, data in enumerate(test_dataloader):
+#     x, y = data
+#     x = x.numpy()
+#     y = y.numpy()
+#     # print(x.shape)
+#     # print(y.shape)
+#     x_test_pred = np.argmax(classifier.predict(x), axis=1)
+#     print(
+#         "Accuracy on benign test samples after adversarial training: %.2f%%"
+#         % (np.sum(x_test_pred == np.argmax(y, axis=1)) / x.shape[0] * 100)
+#     )export https_proxy=http://127.0.0.1:7897 http_proxy=http://127.0.0.1:7897 all_proxy=socks5://127.0.0.1:7897
+
+# trainer.classifier.save('AT-cifar10.pth')
+torch.save(trainer.classifier.model.state_dict(), 'AT-svhn.pth')
+

cifar10.py

@@ -0,0 +1,192 @@
+from PIL import Image
+import numpy as np
+import timm
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import torchvision.transforms as transforms
+from torch.utils.data import Dataset, DataLoader
+from torch.optim.lr_scheduler import MultiStepLR, StepLR
+
+from art.estimators.classification import PyTorchClassifier
+from art.data_generators import PyTorchDataGenerator
+from art.defences.trainer import AdversarialTrainer
+from art.attacks.evasion import ProjectedGradientDescent
+from datasets import load_dataset
+from torchvision.transforms import (CenterCrop,
+                                    Compose,
+                                    Normalize,
+                                    RandomHorizontalFlip,
+                                    RandomResizedCrop,
+                                    Resize,
+                                    ToTensor)
+from tensorflow.keras.utils import to_categorical
+from transformers import ViTImageProcessor
+
+processor = ViTImageProcessor.from_pretrained("google/vit-base-patch16-224-in21k")
+IMAGENET_DEFAULT_MEAN = processor.image_mean
+IMAGENET_DEFAULT_STD = processor.image_std
+
+size = processor.size["height"]
+
+
+"""
+For this example we choose the ResNet18 model as used in the paper (https://proceedings.mlr.press/v97/zhang19p.html)
+The code for the model architecture has been adopted from
+https://github.com/yaodongyu/TRADES/blob/master/models/resnet.py
+"""
+
+
+model = timm.create_model("timm/vit_base_patch16_224.orig_in21k_ft_in1k", pretrained=False)
+model.head = nn.Linear(model.head.in_features, 10)
+model.load_state_dict(
+    torch.hub.load_state_dict_from_url(
+        "https://huggingface.co/edadaltocg/vit_base_patch16_224_in21k_ft_cifar10/resolve/main/pytorch_model.bin",
+        map_location="cuda",
+        file_name="vit_base_patch16_224_in21k_ft_cifar10.pth",
+    )
+)
+
+
+
+# Step 1: Load the CIFAR10 dataset
+train_ds, test_ds = load_dataset('cifar10', split=['train[:5000]', 'test[:2000]'])
+splits = train_ds.train_test_split(test_size=0.1)
+train_ds = splits['train']
+val_ds = splits['test']
+
+train_size=len(train_ds)
+test_size=len(test_ds)
+
+normalize = Normalize(mean=IMAGENET_DEFAULT_MEAN, std=IMAGENET_DEFAULT_STD)
+_train_transforms = Compose(
+        [
+            RandomResizedCrop(size),
+            RandomHorizontalFlip(),
+            ToTensor(),
+            normalize,
+        ]
+    )
+
+_val_transforms = Compose(
+        [
+            Resize(size),
+            CenterCrop(size),
+            ToTensor(),
+            normalize,
+        ]
+    )
+
+def train_transforms(examples):
+    examples['pixel_values'] = [_train_transforms(image.convert("RGB")) for image in examples['img']]
+    return examples
+
+def val_transforms(examples):
+    examples['pixel_values'] = [_val_transforms(image.convert("RGB")) for image in examples['img']]
+    return examples
+
+train_ds.set_transform(train_transforms)
+val_ds.set_transform(val_transforms)
+test_ds.set_transform(val_transforms)
+
+
+def collate_fn(examples):
+    pixel_values = torch.stack([example["pixel_values"] for example in examples])
+    labels = torch.tensor([example["label"] for example in examples])
+    return pixel_values,labels
+
+train_batch_size = 32
+eval_batch_size = 32
+
+def dataset2np(dataset):
+    X = []
+    Y = []
+    for i in range(int(2000)):
+        x,y = dataset[i]["pixel_values"], dataset[i]["label"]
+        y=to_categorical(y,num_classes=10)
+        X.append(x.detach().numpy())
+        Y.append(y)
+    X = np.array(X).astype("float32")
+    Y = np.array(Y).astype("float32")
+    return X,Y
+
+train_dataloader = DataLoader(train_ds, shuffle=True, collate_fn=collate_fn, batch_size=train_batch_size)
+val_dataloader = DataLoader(val_ds, collate_fn=collate_fn, batch_size=eval_batch_size)
+test_dataloader = DataLoader(test_ds, collate_fn=collate_fn, batch_size=eval_batch_size)
+x_test, y_test=dataset2np(test_ds)
+
+
+opt = torch.optim.SGD(model.parameters(), lr=0.1, momentum=0.9, weight_decay=2e-4)
+lr_scheduler = StepLR(opt, step_size=3, gamma=0.1)
+
+criterion = nn.CrossEntropyLoss()
+
+# Step 3: Create the ART classifier
+
+classifier = PyTorchClassifier(
+    model=model,
+    clip_values=(0.0, 1.0),
+    loss=criterion,
+    optimizer=opt,
+    input_shape=(3, size, size),
+    nb_classes=10,
+)
+
+attack = ProjectedGradientDescent(
+    classifier,
+    norm=np.inf,
+    eps=8.0 / 255.0,
+    eps_step=2.0 / 255.0,
+    max_iter=10,
+    targeted=False,
+    num_random_init=1,
+    batch_size=128,
+    verbose=False,
+)
+
+x_test_clean_pred=np.argmax(classifier.predict(x_test), axis=1)
+print(
+    "Accuracy on clean samples before adversarial training: %.2f%%"
+    % (np.sum(x_test_clean_pred == np.argmax(y_test, axis=1)) / x_test.shape[0] * 100)
+)
+
+
+# Step 4: Create the trainer object - AdversarialTrainerTRADESPyTorch
+trainer = AdversarialTrainer(
+    classifier, attack
+)
+
+# Build a Keras image augmentation object and wrap it in ART
+art_datagen = PyTorchDataGenerator(iterator=train_dataloader, size=train_size, batch_size=128)
+
+# Step 5: fit the trainer
+trainer.fit_generator(art_datagen, nb_epochs=50)
+
+
+x_test_pred = np.argmax(classifier.predict(x_test), axis=1)
+print(
+    "Accuracy on benign test samples after adversarial training: %.2f%%"
+    % (np.sum(x_test_pred == np.argmax(y_test, axis=1)) / x_test.shape[0] * 100)
+)
+
+attack_test = ProjectedGradientDescent(
+    classifier,
+    norm=np.inf,
+    eps=8.0 / 255.0,
+    eps_step=2.0 / 255.0,
+    max_iter=20,
+    targeted=False,
+    num_random_init=1,
+    batch_size=128,
+    verbose=False,
+)
+x_test_attack = attack_test.generate(x_test, y=y_test)
+x_test_attack_pred = np.argmax(classifier.predict(x_test_attack), axis=1)
+print(
+    "Accuracy on original PGD adversarial samples after adversarial training: %.2f%%"
+    % (np.sum(x_test_attack_pred == np.argmax(y_test, axis=1)) / x_test.shape[0] * 100)
+)
+torch.save(trainer.classifier.model.state_dict(), 'cifar10_pgd.pth')
+print(
+    "Save the AT model! "
+)

cifar100.py

@@ -0,0 +1,182 @@
+from PIL import Image
+import numpy as np
+import timm
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import torchvision.transforms as transforms
+from torch.utils.data import Dataset, DataLoader
+from torch.optim.lr_scheduler import MultiStepLR, StepLR
+
+from art.estimators.classification import PyTorchClassifier
+from art.data_generators import PyTorchDataGenerator
+from art.defences.trainer import AdversarialTrainer
+from art.attacks.evasion import ProjectedGradientDescent
+from datasets import load_dataset
+from torchvision.transforms import (CenterCrop,
+                                    Compose,
+                                    Normalize,
+                                    RandomHorizontalFlip,
+                                    RandomResizedCrop,
+                                    Resize,
+                                    ToTensor)
+from tensorflow.keras.utils import to_categorical
+from transformers import ViTImageProcessor
+
+processor = ViTImageProcessor.from_pretrained("google/vit-base-patch16-224-in21k")
+IMAGENET_DEFAULT_MEAN = processor.image_mean
+IMAGENET_DEFAULT_STD = processor.image_std
+
+size = processor.size["height"]
+
+
+model = timm.create_model("timm/vit_base_patch16_224.orig_in21k_ft_in1k",
+pretrained=False)
+model.head = nn.Linear(model.head.in_features, 100)
+model.load_state_dict(
+    torch.hub.load_state_dict_from_url(
+        "https://huggingface.co/edadaltocg/vit_base_patch16_224_in21k_ft_cifar100/resolve/main/pytorch_model.bin",
+        map_location="cuda",
+        file_name="vit_base_patch16_224_in21k_ft_cifar100.pth",
+    )
+)
+
+train_ds = load_dataset("uoft-cs/cifar100",split='train')
+test_ds = load_dataset("uoft-cs/cifar100",split='test')
+splits = train_ds.train_test_split(test_size=0.1)
+train_ds = splits['train']
+val_ds = splits['test']
+
+train_size=len(train_ds)
+test_size=len(test_ds)
+
+normalize = Normalize(mean=IMAGENET_DEFAULT_MEAN, std=IMAGENET_DEFAULT_STD)
+_train_transforms = Compose(
+        [
+            RandomResizedCrop(size),
+            RandomHorizontalFlip(),
+            ToTensor(),
+            normalize,
+        ]
+    )
+
+_val_transforms = Compose(
+        [
+            Resize(size),
+            CenterCrop(size),
+            ToTensor(),
+            normalize,
+        ]
+    )
+def train_transforms(examples):
+    examples['pixel_values'] = [_train_transforms(image.convert("RGB")) for image in examples['img']]
+    return examples
+
+def val_transforms(examples):
+    examples['pixel_values'] = [_val_transforms(image.convert("RGB")) for image in examples['img']]
+    return examples
+
+train_ds.set_transform(train_transforms)
+val_ds.set_transform(val_transforms)
+test_ds.set_transform(val_transforms)
+
+def collate_fn(examples):
+    pixel_values = torch.stack([example["pixel_values"] for example in examples])
+    labels = torch.tensor([example["fine_label"] for example in examples])
+    return pixel_values,labels
+
+train_batch_size = 64
+eval_batch_size = 64
+
+def dataset2np(dataset):
+    X = []
+    Y = []
+    for i in range(int(2000)):
+        x,y = dataset[i]["pixel_values"], dataset[i]["fine_label"]
+        y=to_categorical(y,num_classes=100)
+        X.append(x.detach().numpy())
+        Y.append(y)
+    X = np.array(X).astype("float32")
+    Y = np.array(Y).astype("float32")
+    return X,Y
+
+train_dataloader = DataLoader(train_ds, shuffle=True, collate_fn=collate_fn, batch_size=train_batch_size)
+val_dataloader = DataLoader(val_ds, collate_fn=collate_fn, batch_size=eval_batch_size)
+test_dataloader = DataLoader(test_ds, collate_fn=collate_fn, batch_size=eval_batch_size)
+x_test, y_test=dataset2np(test_ds)
+
+
+opt = torch.optim.SGD(model.parameters(), lr=0.1, momentum=0.9, weight_decay=2e-4)
+lr_scheduler = StepLR(opt, step_size=3, gamma=0.1)
+
+criterion = nn.CrossEntropyLoss()
+
+# Step 3: Create the ART classifier
+
+classifier = PyTorchClassifier(
+    model=model,
+    clip_values=(0.0, 1.0),
+    loss=criterion,
+    optimizer=opt,
+    input_shape=(3, size, size),
+    nb_classes=100,
+)
+
+attack = ProjectedGradientDescent(
+    classifier,
+    norm=np.inf,
+    eps=8.0 / 255.0,
+    eps_step=2.0 / 255.0,
+    max_iter=10,
+    targeted=False,
+    num_random_init=1,
+    batch_size=128,
+    verbose=False,
+)
+
+x_test_clean_pred=np.argmax(classifier.predict(x_test), axis=1)
+print(
+    "Accuracy on clean samples before adversarial training: %.2f%%"
+    % (np.sum(x_test_clean_pred == np.argmax(y_test, axis=1)) / x_test.shape[0] * 100)
+)
+
+
+# Step 4: Create the trainer object - AdversarialTrainerTRADESPyTorch
+trainer = AdversarialTrainer(
+    classifier, attack
+)
+
+# Build a Keras image augmentation object and wrap it in ART
+art_datagen = PyTorchDataGenerator(iterator=train_dataloader, size=train_size, batch_size=128)
+
+# Step 5: fit the trainer
+trainer.fit_generator(art_datagen, nb_epochs=50)
+
+
+x_test_pred = np.argmax(classifier.predict(x_test), axis=1)
+print(
+    "Accuracy on benign test samples after adversarial training: %.2f%%"
+    % (np.sum(x_test_pred == np.argmax(y_test, axis=1)) / x_test.shape[0] * 100)
+)
+
+attack_test = ProjectedGradientDescent(
+    classifier,
+    norm=np.inf,
+    eps=8.0 / 255.0,
+    eps_step=2.0 / 255.0,
+    max_iter=20,
+    targeted=False,
+    num_random_init=1,
+    batch_size=128,
+    verbose=False,
+)
+x_test_attack = attack_test.generate(x_test, y=y_test)
+x_test_attack_pred = np.argmax(classifier.predict(x_test_attack), axis=1)
+print(
+    "Accuracy on original PGD adversarial samples after adversarial training: %.2f%%"
+    % (np.sum(x_test_attack_pred == np.argmax(y_test, axis=1)) / x_test.shape[0] * 100)
+)
+torch.save(trainer.classifier.model.state_dict(), 'cifar100-pgd.pth')
+print(
+    "Save the AT model! "
+)

dataset/base.py

@@ -8,8 +8,8 @@ import torch
 import random
 from PIL import Image
 from torchvision.transforms import Compose, Resize, CenterCrop, ToTensor, Normalize
-from model.simple_tokenizer import SimpleTokenizer as Tokenizer
-
+# from model.simple_tokenizer import SimpleTokenizer as Tokenizer
+from transformers import AutoTokenizer
 
 class BaseDataset(Dataset):
 
@@ -19,7 +19,7 @@ class BaseDataset(Dataset):
             indexs: dict,
             labels: dict,
             is_train=True,
-            tokenizer=Tokenizer(),
+            tokenizer=AutoTokenizer.from_pretrained("openai/clip-vit-base-patch32", model_max_length=77, truncation=True),
             maxWords=32,
             imageResolution=224,
             npy=False):

dataset/dataloader.py

@@ -10,7 +10,7 @@ def split_data(captions, indexs, labels, query_num=5000, train_num=10000, seed=N
     random_index = np.random.permutation(range(len(indexs)))
     query_index = random_index[: query_num]
     train_index = random_index[query_num: query_num + train_num]
-    retrieval_index = random_index[query_num:]
+    retrieval_index = random_index[query_num:-100000]
 
     query_indexs = indexs[query_index]
     query_captions = captions[query_index]

dataset/make_nuswide.py

@@ -1,107 +0,0 @@
-import os
-import scipy.io as scio
-import numpy as np
-
-# mkdir mat
-# mv make_nuswide.py mat
-# python make_nuswide.py
-root_dir = "PATH/TO/YOUR/DOWNLOAD/DIR/"
-
-
-imageListFile = os.path.join(root_dir, "/Low-Level-Features/ImageList/Imagelist.txt")
-labelPath = os.path.join(root_dir, "/nuswide/Groundtruth/AllLabels")
-textFile = os.path.join(root_dir, "/Low-Level-Features/NUS_WID_Tags/All_Tags.txt")
-classIndexFile = os.path.join(root_dir, "/Low-Level-Features/Concepts81.txt")
-
-# you can use the image urls to download images
-imagePath = os.path.join(root_dir, "nuswide/Flickr")
-
-with open(imageListFile, "r") as f:
-    indexs = f.readlines()
-
-indexs = [os.path.join(imagePath, item.strip().replace("\\", "/")) for item in indexs]
-print("indexs length:", len(indexs))
-
-#class_index = {}
-#with open(classIndexFile, "r") as f:
-#    data = f.readlines()
-#
-#for i, item in enumerate(data):
-#    class_index.update({item.strip(): i})
-
-captions = []
-with open(textFile, "r") as f:
-    for line in f:
-        if len(line.strip()) == 0:
-            print("some line empty!")
-            continue
-        caption = line.split()[1:]
-        caption = " ".join(caption).strip()
-        if len(caption) == 0:
-             caption = "123456"
-        captions.append(caption)
-
-print("captions length:", len(captions))
-
-#labels = np.zeros([len(indexs), len(class_index)], dtype=np.int8)
-# label_lists = os.listdir(labelPath)
-with open(os.path.join(root_dir, "/nuswide/Groundtruth/used_label.txt")) as f:
-    label_lists = f.readlines()
-label_lists = [item.strip() for item in label_lists]
-
-class_index = {}
-for i, item in enumerate(label_lists):
-    class_index.update({item: i})
-
-labels = np.zeros([len(indexs), len(class_index)], dtype=np.int8)
-
-for item in label_lists:
-    path = os.path.join(labelPath, item)
-    class_label = item# .split(".")[0].split("_")[-1]
-
-    with open(path, "r") as f:
-        data = f.readlines()
-    for i, val in enumerate(data):
-        labels[i][class_index[class_label]] = 1 if val.strip() == "1" else 0
-print("labels sum:", labels.sum())
-
-not_used_id = []
-with open(os.path.join(root_dir, "/nuswide/Groundtruth/not_used_id.txt")) as f:
-    not_used_id = f.readlines()
-not_used_id = [int(item.strip()) for item in not_used_id]
-
-# for item in not_used_id:
-#     indexs.pop(item)
-#     captions.pop(item)
-#     labels = np.delete(labels, item, 0)
-ind = list(range(len(indexs)))
-for item in not_used_id:
-    ind.remove(item)
-    indexs[item] = ""
-    captions[item] = ""
-indexs = [item for item in indexs if item != ""]
-captions = [item for item in captions if item != ""]
-ind = np.asarray(ind)
-labels = labels[ind]
-# ind = range(len(indexs))
-
-print("indexs length:", len(indexs))
-print("captions length:", len(captions))
-print("labels shape:", labels.shape)
-
-indexs = {"index": indexs}
-captions = {"caption": captions}
-labels = {"category": labels}
-
-scio.savemat(os.path.join(root_dir, "/mat/index.mat"), indexs)
-# scio.savemat("caption.mat", captions)
-scio.savemat(os.path.join(root_dir, "/mat/label.mat"), labels)
-
-
-captions = [item + "\n" for item in captions["caption"]]
-
-with open(os.path.join(root_dir, "/mat/caption.txt"), "w") as f:
-    f.writelines(captions)
-
-print("finished!")
-

main.py

@@ -1,5 +1,5 @@
-from train.hash_train import Trainer
-
+from train.text_train import Trainer
+# from train.hash_train import Trainer
 
 if __name__ == "__main__":
 

mlp_mix.py

@@ -0,0 +1,92 @@
+import typing as tp
+
+import jax
+import jax.numpy as jnp
+from einops import rearrange
+from functools import partial
+
+
+
+
+from flax import nnx
+
+class FeedForward(nnx.Module):
+  def __init__(self, dim, hidden_dim, dropout , rngs: nnx.Rngs):
+    self.net=nnx.Sequential(
+        nnx.Linear(dim, hidden_dim , rngs=rngs),
+            partial(nnx.gelu),
+            nnx.Dropout(dropout , rngs=rngs),
+            nnx.Linear(hidden_dim, dim , rngs=rngs),
+            nnx.Dropout(dropout , rngs=rngs)
+    )
+    
+
+  def __call__(self, x):
+    return self.net(x)
+
+
+class MixerBlock(nnx.Module):
+
+    def __init__(self, dim, num_patch, token_dim, channel_dim, dropout , rngs: nnx.Rngs):
+        super().__init__()
+        self.ln1=nnx.LayerNorm(dim, rngs=rngs)
+        self.ffn1=FeedForward(num_patch,token_dim,dropout,rngs=rngs)
+        
+        self.ln2=nnx.LayerNorm(dim, rngs=rngs)
+        self.ffn2=FeedForward(dim, channel_dim, dropout, rngs=rngs)
+
+        
+
+    def __call__(self, x):
+        # print(x.shape)
+        x = x + self.ffn1(self.ln1(x))
+
+        x = x + self.ffn2(self.ln2(x))
+
+        return x
+    
+class MLPMixer(nnx.Module):
+
+    def __init__(self, in_channels, dim, num_classes, patch_size,dropout, image_size, depth, token_dim, channel_dim, rngs: nnx.Rngs):
+        super().__init__()
+
+        assert image_size % patch_size == 0, 'Image dimensions must be divisible by the patch size.'
+        self.num_patch =  (image_size// patch_size) ** 2
+        
+        self.to_patch_embedding = nnx.Sequential(
+            nnx.Conv(in_channels, dim, kernel_size=(patch_size, patch_size), rngs=rngs),
+        )
+        self.mixer_blocks=[]
+
+        for _ in range(depth):
+            self.mixer_blocks.append(MixerBlock(dim, self.num_patch, token_dim, channel_dim,dropout, rngs=rngs))
+
+        self.layer_norm = nnx.LayerNorm(dim, rngs=rngs)
+
+        self.mlp_head = nnx.Sequential(
+            nnx.Linear(dim, num_classes, rngs=rngs)
+        )
+
+    def __call__(self, x):
+
+
+        x = self.to_patch_embedding(x)
+
+        for mixer_block in self.mixer_blocks:
+            x = mixer_block(x)
+
+        x = self.layer_norm(x)
+
+        x = jnp.mean(x, axis=1)
+
+        return self.mlp_head(x)
+    
+if __name__ == "__main__":
+    img = jnp.ones([1, 3, 224, 224])
+
+    model = MLPMixer(in_channels=3, image_size=224, patch_size=16,dropout=0.2, num_classes=1000,
+                     dim=512, depth=8, token_dim=256, channel_dim=2048,rngs=nnx.Rngs(0))
+    # nnx.display(model)
+    out_img = model(jnp.ones((1,  224, 224,3)))
+
+    print("Shape of out :", out_img.shape)  # [B, in_channels, image_size, image_size]

model/bert_tokenizer.py

@@ -0,0 +1,543 @@
+# coding=utf-8
+# Copyright 2018 The Google AI Language Team Authors and The HuggingFace Inc. team.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""Tokenization classes for Bert."""
+
+
+import collections
+import os
+import unicodedata
+from typing import List, Optional, Tuple
+
+from transformers.tokenization_utils import PreTrainedTokenizer, _is_control, _is_punctuation, _is_whitespace
+from transformers.utils import logging
+
+
+logger = logging.get_logger(__name__)
+
+VOCAB_FILES_NAMES = {"vocab_file": "vocab.txt"}
+
+PRETRAINED_VOCAB_FILES_MAP = {
+    "vocab_file": {
+        "bert-base-uncased": "https://huggingface.co/bert-base-uncased/resolve/main/vocab.txt",
+        "bert-large-uncased": "https://huggingface.co/bert-large-uncased/resolve/main/vocab.txt",
+        "bert-base-cased": "https://huggingface.co/bert-base-cased/resolve/main/vocab.txt",
+        "bert-large-cased": "https://huggingface.co/bert-large-cased/resolve/main/vocab.txt",
+        "bert-base-multilingual-uncased": "https://huggingface.co/bert-base-multilingual-uncased/resolve/main/vocab.txt",
+        "bert-base-multilingual-cased": "https://huggingface.co/bert-base-multilingual-cased/resolve/main/vocab.txt",
+        "bert-base-chinese": "https://huggingface.co/bert-base-chinese/resolve/main/vocab.txt",
+        "bert-base-german-cased": "https://huggingface.co/bert-base-german-cased/resolve/main/vocab.txt",
+        "bert-large-uncased-whole-word-masking": "https://huggingface.co/bert-large-uncased-whole-word-masking/resolve/main/vocab.txt",
+        "bert-large-cased-whole-word-masking": "https://huggingface.co/bert-large-cased-whole-word-masking/resolve/main/vocab.txt",
+        "bert-large-uncased-whole-word-masking-finetuned-squad": "https://huggingface.co/bert-large-uncased-whole-word-masking-finetuned-squad/resolve/main/vocab.txt",
+        "bert-large-cased-whole-word-masking-finetuned-squad": "https://huggingface.co/bert-large-cased-whole-word-masking-finetuned-squad/resolve/main/vocab.txt",
+        "bert-base-cased-finetuned-mrpc": "https://huggingface.co/bert-base-cased-finetuned-mrpc/resolve/main/vocab.txt",
+        "bert-base-german-dbmdz-cased": "https://huggingface.co/bert-base-german-dbmdz-cased/resolve/main/vocab.txt",
+        "bert-base-german-dbmdz-uncased": "https://huggingface.co/bert-base-german-dbmdz-uncased/resolve/main/vocab.txt",
+        "TurkuNLP/bert-base-finnish-cased-v1": "https://huggingface.co/TurkuNLP/bert-base-finnish-cased-v1/resolve/main/vocab.txt",
+        "TurkuNLP/bert-base-finnish-uncased-v1": "https://huggingface.co/TurkuNLP/bert-base-finnish-uncased-v1/resolve/main/vocab.txt",
+        "wietsedv/bert-base-dutch-cased": "https://huggingface.co/wietsedv/bert-base-dutch-cased/resolve/main/vocab.txt",
+    }
+}
+
+PRETRAINED_POSITIONAL_EMBEDDINGS_SIZES = {
+    "bert-base-uncased": 512,
+    "bert-large-uncased": 512,
+    "bert-base-cased": 512,
+    "bert-large-cased": 512,
+    "bert-base-multilingual-uncased": 512,
+    "bert-base-multilingual-cased": 512,
+    "bert-base-chinese": 512,
+    "bert-base-german-cased": 512,
+    "bert-large-uncased-whole-word-masking": 512,
+    "bert-large-cased-whole-word-masking": 512,
+    "bert-large-uncased-whole-word-masking-finetuned-squad": 512,
+    "bert-large-cased-whole-word-masking-finetuned-squad": 512,
+    "bert-base-cased-finetuned-mrpc": 512,
+    "bert-base-german-dbmdz-cased": 512,
+    "bert-base-german-dbmdz-uncased": 512,
+    "TurkuNLP/bert-base-finnish-cased-v1": 512,
+    "TurkuNLP/bert-base-finnish-uncased-v1": 512,
+    "wietsedv/bert-base-dutch-cased": 512,
+}
+
+PRETRAINED_INIT_CONFIGURATION = {
+    "bert-base-uncased": {"do_lower_case": True},
+    "bert-large-uncased": {"do_lower_case": True},
+    "bert-base-cased": {"do_lower_case": False},
+    "bert-large-cased": {"do_lower_case": False},
+    "bert-base-multilingual-uncased": {"do_lower_case": True},
+    "bert-base-multilingual-cased": {"do_lower_case": False},
+    "bert-base-chinese": {"do_lower_case": False},
+    "bert-base-german-cased": {"do_lower_case": False},
+    "bert-large-uncased-whole-word-masking": {"do_lower_case": True},
+    "bert-large-cased-whole-word-masking": {"do_lower_case": False},
+    "bert-large-uncased-whole-word-masking-finetuned-squad": {"do_lower_case": True},
+    "bert-large-cased-whole-word-masking-finetuned-squad": {"do_lower_case": False},
+    "bert-base-cased-finetuned-mrpc": {"do_lower_case": False},
+    "bert-base-german-dbmdz-cased": {"do_lower_case": False},
+    "bert-base-german-dbmdz-uncased": {"do_lower_case": True},
+    "TurkuNLP/bert-base-finnish-cased-v1": {"do_lower_case": False},
+    "TurkuNLP/bert-base-finnish-uncased-v1": {"do_lower_case": True},
+    "wietsedv/bert-base-dutch-cased": {"do_lower_case": False},
+}
+
+
+def load_vocab(vocab_file):
+    """Loads a vocabulary file into a dictionary."""
+    vocab = collections.OrderedDict()
+    with open(vocab_file, "r", encoding="utf-8") as reader:
+        tokens = reader.readlines()
+    for index, token in enumerate(tokens):
+        token = token.rstrip("\n")
+        vocab[token] = index
+    return vocab
+
+
+def whitespace_tokenize(text):
+    """Runs basic whitespace cleaning and splitting on a piece of text."""
+    text = text.strip()
+    if not text:
+        return []
+    tokens = text.split()
+    return tokens
+
+
+class BertTokenizer(PreTrainedTokenizer):
+    r"""
+    Construct a BERT tokenizer. Based on WordPiece.
+    This tokenizer inherits from :class:`~transformers.PreTrainedTokenizer` which contains most of the main methods.
+    Users should refer to this superclass for more information regarding those methods.
+    Args:
+        vocab_file (:obj:`str`):
+            File containing the vocabulary.
+        do_lower_case (:obj:`bool`, `optional`, defaults to :obj:`True`):
+            Whether or not to lowercase the input when tokenizing.
+        do_basic_tokenize (:obj:`bool`, `optional`, defaults to :obj:`True`):
+            Whether or not to do basic tokenization before WordPiece.
+        never_split (:obj:`Iterable`, `optional`):
+            Collection of tokens which will never be split during tokenization. Only has an effect when
+            :obj:`do_basic_tokenize=True`
+        unk_token (:obj:`str`, `optional`, defaults to :obj:`"[UNK]"`):
+            The unknown token. A token that is not in the vocabulary cannot be converted to an ID and is set to be this
+            token instead.
+        sep_token (:obj:`str`, `optional`, defaults to :obj:`"[SEP]"`):
+            The separator token, which is used when building a sequence from multiple sequences, e.g. two sequences for
+            sequence classification or for a text and a question for question answering. It is also used as the last
+            token of a sequence built with special tokens.
+        pad_token (:obj:`str`, `optional`, defaults to :obj:`"[PAD]"`):
+            The token used for padding, for example when batching sequences of different lengths.
+        cls_token (:obj:`str`, `optional`, defaults to :obj:`"[CLS]"`):
+            The classifier token which is used when doing sequence classification (classification of the whole sequence
+            instead of per-token classification). It is the first token of the sequence when built with special tokens.
+        mask_token (:obj:`str`, `optional`, defaults to :obj:`"[MASK]"`):
+            The token used for masking values. This is the token used when training this model with masked language
+            modeling. This is the token which the model will try to predict.
+        tokenize_chinese_chars (:obj:`bool`, `optional`, defaults to :obj:`True`):
+            Whether or not to tokenize Chinese characters.
+            This should likely be deactivated for Japanese (see this `issue
+            <https://github.com/huggingface/transformers/issues/328>`__).
+        strip_accents: (:obj:`bool`, `optional`):
+            Whether or not to strip all accents. If this option is not specified, then it will be determined by the
+            value for :obj:`lowercase` (as in the original BERT).
+    """
+
+    vocab_files_names = VOCAB_FILES_NAMES
+    pretrained_vocab_files_map = PRETRAINED_VOCAB_FILES_MAP
+    pretrained_init_configuration = PRETRAINED_INIT_CONFIGURATION
+    max_model_input_sizes = PRETRAINED_POSITIONAL_EMBEDDINGS_SIZES
+
+    def __init__(
+        self,
+        vocab_file,
+        do_lower_case=True,
+        do_basic_tokenize=True,
+        never_split=None,
+        unk_token="[UNK]",
+        sep_token="[SEP]",
+        pad_token="[PAD]",
+        cls_token="[CLS]",
+        mask_token="[MASK]",
+        tokenize_chinese_chars=True,
+        strip_accents=None,
+        **kwargs
+    ):
+        super().__init__(
+            do_lower_case=do_lower_case,
+            do_basic_tokenize=do_basic_tokenize,
+            never_split=never_split,
+            unk_token=unk_token,
+            sep_token=sep_token,
+            pad_token=pad_token,
+            cls_token=cls_token,
+            mask_token=mask_token,
+            tokenize_chinese_chars=tokenize_chinese_chars,
+            strip_accents=strip_accents,
+            **kwargs,
+        )
+
+        if not os.path.isfile(vocab_file):
+            raise ValueError(
+                "Can't find a vocabulary file at path '{}'. To load the vocabulary from a Google pretrained "
+                "model use `tokenizer = BertTokenizer.from_pretrained(PRETRAINED_MODEL_NAME)`".format(vocab_file)
+            )
+        self.vocab = load_vocab(vocab_file)
+        self.ids_to_tokens = collections.OrderedDict([(ids, tok) for tok, ids in self.vocab.items()])
+        self.do_basic_tokenize = do_basic_tokenize
+        if do_basic_tokenize:
+            self.basic_tokenizer = BasicTokenizer(
+                do_lower_case=do_lower_case,
+                never_split=never_split,
+                tokenize_chinese_chars=tokenize_chinese_chars,
+                strip_accents=strip_accents,
+            )
+        self.wordpiece_tokenizer = WordpieceTokenizer(vocab=self.vocab, unk_token=self.unk_token)
+
+    @property
+    def do_lower_case(self):
+        return self.basic_tokenizer.do_lower_case
+
+    @property
+    def vocab_size(self):
+        return len(self.vocab)
+
+    def get_vocab(self):
+        return dict(self.vocab, **self.added_tokens_encoder)
+
+    def _tokenize(self, text):
+        split_tokens = []
+        if self.do_basic_tokenize:
+            for token in self.basic_tokenizer.tokenize(text, never_split=self.all_special_tokens):
+
+                # If the token is part of the never_split set
+                if token in self.basic_tokenizer.never_split:
+                    split_tokens.append(token)
+                else:
+                    split_tokens += self.wordpiece_tokenizer.tokenize(token)
+        else:
+            split_tokens = self.wordpiece_tokenizer.tokenize(text)
+        return split_tokens
+
+    def _convert_token_to_id(self, token):
+        """ Converts a token (str) in an id using the vocab. """
+        return self.vocab.get(token, self.vocab.get(self.unk_token))
+    
+    # def _convert_tokens_to_ids(self, tokens):
+    #     """ Converts a token (str) in an id using the vocab. """
+    #     return [self._convert_token_to_id(token) for token in tokens]
+
+    def _convert_id_to_token(self, index):
+        """Converts an index (integer) in a token (str) using the vocab."""
+        return self.ids_to_tokens.get(index, self.unk_token)
+
+    def convert_tokens_to_string(self, tokens):
+        """ Converts a sequence of tokens (string) in a single string. """
+        out_string = " ".join(tokens).replace(" ##", "").strip()
+        return out_string
+
+    def build_inputs_with_special_tokens(
+        self, token_ids_0: List[int], token_ids_1: Optional[List[int]] = None
+    ) -> List[int]:
+        """
+        Build model inputs from a sequence or a pair of sequence for sequence classification tasks by concatenating and
+        adding special tokens. A BERT sequence has the following format:
+        - single sequence: ``[CLS] X ``
+        - pair of sequences: ``[CLS] A [SEP] B [SEP]``
+        Args:
+            token_ids_0 (:obj:`List[int]`):
+                List of IDs to which the special tokens will be added.
+            token_ids_1 (:obj:`List[int]`, `optional`):
+                Optional second list of IDs for sequence pairs.
+        Returns:
+            :obj:`List[int]`: List of `input IDs <../glossary.html#input-ids>`__ with the appropriate special tokens.
+        """
+        if token_ids_1 is None:
+            return [self.cls_token_id] + token_ids_0
+        cls = [self.cls_token_id]
+        sep = [self.sep_token_id]
+        return cls + token_ids_0 + sep + token_ids_1 + sep
+
+    def get_special_tokens_mask(
+        self, token_ids_0: List[int], token_ids_1: Optional[List[int]] = None, already_has_special_tokens: bool = False
+    ) -> List[int]:
+        """
+        Retrieve sequence ids from a token list that has no special tokens added. This method is called when adding
+        special tokens using the tokenizer ``prepare_for_model`` method.
+        Args:
+            token_ids_0 (:obj:`List[int]`):
+                List of IDs.
+            token_ids_1 (:obj:`List[int]`, `optional`):
+                Optional second list of IDs for sequence pairs.
+            already_has_special_tokens (:obj:`bool`, `optional`, defaults to :obj:`False`):
+                Whether or not the token list is already formatted with special tokens for the model.
+        Returns:
+            :obj:`List[int]`: A list of integers in the range [0, 1]: 1 for a special token, 0 for a sequence token.
+        """
+
+        if already_has_special_tokens:
+            if token_ids_1 is not None:
+                raise ValueError(
+                    "You should not supply a second sequence if the provided sequence of "
+                    "ids is already formatted with special tokens for the model."
+                )
+            return list(map(lambda x: 1 if x in [self.sep_token_id, self.cls_token_id] else 0, token_ids_0))
+
+        if token_ids_1 is not None:
+            return [1] + ([0] * len(token_ids_0)) + [1] + ([0] * len(token_ids_1)) + [1]
+        return [1] + ([0] * len(token_ids_0)) + [1]
+
+    def create_token_type_ids_from_sequences(
+        self, token_ids_0: List[int], token_ids_1: Optional[List[int]] = None
+    ) -> List[int]:
+        """
+        Create a mask from the two sequences passed to be used in a sequence-pair classification task. A BERT sequence
+        pair mask has the following format:
+        ::
+            0 0 0 0 0 0 0 0 0 0 0 1 1 1 1 1 1 1 1 1
+            | first sequence    | second sequence |
+        If :obj:`token_ids_1` is :obj:`None`, this method only returns the first portion of the mask (0s).
+        Args:
+            token_ids_0 (:obj:`List[int]`):
+                List of IDs.
+            token_ids_1 (:obj:`List[int]`, `optional`):
+                Optional second list of IDs for sequence pairs.
+        Returns:
+            :obj:`List[int]`: List of `token type IDs <../glossary.html#token-type-ids>`_ according to the given
+            sequence(s).
+        """
+        sep = [self.sep_token_id]
+        cls = [self.cls_token_id]
+        if token_ids_1 is None:
+            return len(cls + token_ids_0 + sep) * [0]
+        return len(cls + token_ids_0 + sep) * [0] + len(token_ids_1 + sep) * [1]
+
+    def save_vocabulary(self, save_directory: str, filename_prefix: Optional[str] = None) -> Tuple[str]:
+        index = 0
+        if os.path.isdir(save_directory):
+            vocab_file = os.path.join(
+                save_directory, (filename_prefix + "-" if filename_prefix else "") + VOCAB_FILES_NAMES["vocab_file"]
+            )
+        else:
+            vocab_file = (filename_prefix + "-" if filename_prefix else "") + save_directory
+        with open(vocab_file, "w", encoding="utf-8") as writer:
+            for token, token_index in sorted(self.vocab.items(), key=lambda kv: kv[1]):
+                if index != token_index:
+                    logger.warning(
+                        "Saving vocabulary to {}: vocabulary indices are not consecutive."
+                        " Please check that the vocabulary is not corrupted!".format(vocab_file)
+                    )
+                    index = token_index
+                writer.write(token + "\n")
+                index += 1
+        return (vocab_file,)
+
+
+class BasicTokenizer(object):
+    """
+    Constructs a BasicTokenizer that will run basic tokenization (punctuation splitting, lower casing, etc.).
+    Args:
+        do_lower_case (:obj:`bool`, `optional`, defaults to :obj:`True`):
+            Whether or not to lowercase the input when tokenizing.
+        never_split (:obj:`Iterable`, `optional`):
+            Collection of tokens which will never be split during tokenization. Only has an effect when
+            :obj:`do_basic_tokenize=True`
+        tokenize_chinese_chars (:obj:`bool`, `optional`, defaults to :obj:`True`):
+            Whether or not to tokenize Chinese characters.
+            This should likely be deactivated for Japanese (see this `issue
+            <https://github.com/huggingface/transformers/issues/328>`__).
+        strip_accents: (:obj:`bool`, `optional`):
+            Whether or not to strip all accents. If this option is not specified, then it will be determined by the
+            value for :obj:`lowercase` (as in the original BERT).
+    """
+
+    def __init__(self, do_lower_case=True, never_split=None, tokenize_chinese_chars=True, strip_accents=None):
+        if never_split is None:
+            never_split = []
+        self.do_lower_case = do_lower_case
+        self.never_split = set(never_split)
+        self.tokenize_chinese_chars = tokenize_chinese_chars
+        self.strip_accents = strip_accents
+
+    def tokenize(self, text, never_split=None):
+        """
+        Basic Tokenization of a piece of text. Split on "white spaces" only, for sub-word tokenization, see
+        WordPieceTokenizer.
+        Args:
+            **never_split**: (`optional`) list of str
+                Kept for backward compatibility purposes. Now implemented directly at the base class level (see
+                :func:`PreTrainedTokenizer.tokenize`) List of token not to split.
+        """
+        # union() returns a new set by concatenating the two sets.
+        never_split = self.never_split.union(set(never_split)) if never_split else self.never_split
+        text = self._clean_text(text)
+
+        # This was added on November 1st, 2018 for the multilingual and Chinese
+        # models. This is also applied to the English models now, but it doesn't
+        # matter since the English models were not trained on any Chinese data
+        # and generally don't have any Chinese data in them (there are Chinese
+        # characters in the vocabulary because Wikipedia does have some Chinese
+        # words in the English Wikipedia.).
+        if self.tokenize_chinese_chars:
+            text = self._tokenize_chinese_chars(text)
+        orig_tokens = whitespace_tokenize(text)
+        split_tokens = []
+        for token in orig_tokens:
+            if token not in never_split:
+                if self.do_lower_case:
+                    token = token.lower()
+                    if self.strip_accents is not False:
+                        token = self._run_strip_accents(token)
+                elif self.strip_accents:
+                    token = self._run_strip_accents(token)
+            split_tokens.extend(self._run_split_on_punc(token, never_split))
+
+        output_tokens = whitespace_tokenize(" ".join(split_tokens))
+        return output_tokens
+
+    def _run_strip_accents(self, text):
+        """Strips accents from a piece of text."""
+        text = unicodedata.normalize("NFD", text)
+        output = []
+        for char in text:
+            cat = unicodedata.category(char)
+            if cat == "Mn":
+                continue
+            output.append(char)
+        return "".join(output)
+
+    def _run_split_on_punc(self, text, never_split=None):
+        """Splits punctuation on a piece of text."""
+        if never_split is not None and text in never_split:
+            return [text]
+        chars = list(text)
+        i = 0
+        start_new_word = True
+        output = []
+        while i < len(chars):
+            char = chars[i]
+            if _is_punctuation(char):
+                output.append([char])
+                start_new_word = True
+            else:
+                if start_new_word:
+                    output.append([])
+                start_new_word = False
+                output[-1].append(char)
+            i += 1
+
+        return ["".join(x) for x in output]
+
+    def _tokenize_chinese_chars(self, text):
+        """Adds whitespace around any CJK character."""
+        output = []
+        for char in text:
+            cp = ord(char)
+            if self._is_chinese_char(cp):
+                output.append(" ")
+                output.append(char)
+                output.append(" ")
+            else:
+                output.append(char)
+        return "".join(output)
+
+    def _is_chinese_char(self, cp):
+        """Checks whether CP is the codepoint of a CJK character."""
+        # This defines a "chinese character" as anything in the CJK Unicode block:
+        #   https://en.wikipedia.org/wiki/CJK_Unified_Ideographs_(Unicode_block)
+        #
+        # Note that the CJK Unicode block is NOT all Japanese and Korean characters,
+        # despite its name. The modern Korean Hangul alphabet is a different block,
+        # as is Japanese Hiragana and Katakana. Those alphabets are used to write
+        # space-separated words, so they are not treated specially and handled
+        # like the all of the other languages.
+        if (
+            (cp >= 0x4E00 and cp <= 0x9FFF)
+            or (cp >= 0x3400 and cp <= 0x4DBF)  #
+            or (cp >= 0x20000 and cp <= 0x2A6DF)  #
+            or (cp >= 0x2A700 and cp <= 0x2B73F)  #
+            or (cp >= 0x2B740 and cp <= 0x2B81F)  #
+            or (cp >= 0x2B820 and cp <= 0x2CEAF)  #
+            or (cp >= 0xF900 and cp <= 0xFAFF)
+            or (cp >= 0x2F800 and cp <= 0x2FA1F)  #
+        ):  #
+            return True
+
+        return False
+
+    def _clean_text(self, text):
+        """Performs invalid character removal and whitespace cleanup on text."""
+        output = []
+        for char in text:
+            cp = ord(char)
+            if cp == 0 or cp == 0xFFFD or _is_control(char):
+                continue
+            if _is_whitespace(char):
+                output.append(" ")
+            else:
+                output.append(char)
+        return "".join(output)
+
+
+class WordpieceTokenizer(object):
+    """Runs WordPiece tokenization."""
+
+    def __init__(self, vocab, unk_token, max_input_chars_per_word=100):
+        self.vocab = vocab
+        self.unk_token = unk_token
+        self.max_input_chars_per_word = max_input_chars_per_word
+
+    def tokenize(self, text):
+        """
+        Tokenizes a piece of text into its word pieces. This uses a greedy longest-match-first algorithm to perform
+        tokenization using the given vocabulary.
+        For example, :obj:`input = "unaffable"` wil return as output :obj:`["un", "##aff", "##able"]`.
+        Args:
+          text: A single token or whitespace separated tokens. This should have
+            already been passed through `BasicTokenizer`.
+        Returns:
+          A list of wordpiece tokens.
+        """
+
+        output_tokens = []
+        for token in whitespace_tokenize(text):
+            chars = list(token)
+            if len(chars) > self.max_input_chars_per_word:
+                output_tokens.append(self.unk_token)
+                continue
+
+            is_bad = False
+            start = 0
+            sub_tokens = []
+            while start < len(chars):
+                end = len(chars)
+                cur_substr = None
+                while start < end:
+                    substr = "".join(chars[start:end])
+                    if start > 0:
+                        substr = "##" + substr
+                    if substr in self.vocab:
+                        cur_substr = substr
+                        break
+                    end -= 1
+                if cur_substr is None:
+                    is_bad = True
+                    break
+                sub_tokens.append(cur_substr)
+                start = end
+
+            if is_bad:
+                output_tokens.append(self.unk_token)
+            else:
+                output_tokens.extend(sub_tokens)
+        return output_tokens

model/simple_tokenizer.py

@@ -130,6 +130,17 @@ class SimpleTokenizer(object):
         text = ''.join([self.decoder[token] for token in tokens])
         text = bytearray([self.byte_decoder[c] for c in text]).decode('utf-8', errors="replace").replace('</w>', ' ')
         return text
+
+    # def my_decode(self, tokens):
+    #     print(tokens)
+    #     tem_token=[]
+    #     for i in tokens:
+    #         if i in self.decoder.keys():
+    #             tem_token.append(self.decoder[i])
+    #     print(tem_token)
+    #     text = ''.join(tem_token)
+    #     text = bytearray([self.byte_decoder[c] for c in text]).decode('utf-8', errors="replace").replace('</w>', ' ')
+    #     return text
     
     def tokenize(self, text):
         tokens = []
@@ -141,3 +152,6 @@ class SimpleTokenizer(object):
 
     def convert_tokens_to_ids(self, tokens):
         return [self.encoder[bpe_token] for bpe_token in tokens]
+
+    def convert_ids_to_tokens(self, ids):
+        return [self.decoder[id] for id in ids]

run.sh

@@ -0,0 +1,10 @@
+#!/bin/bash
+set -e
+
+
+export https_proxy=http://127.0.0.1:7897 http_proxy=http://127.0.0.1:7897 all_proxy=socks5://127.0.0.1:7897
+# CUDA_VISIBLE_DEVICES=0 python main.py --method CSQ --bit 32
+# CUDA_VISIBLE_DEVICES=0 python main.py --method CSQ --bit 64
+CUDA_VISIBLE_DEVICES=0 python main.py --victim ViT-B/16 --output-dim 512
+CUDA_VISIBLE_DEVICES=0 python main.py --victim ViT-B/32 --output-dim 512
+CUDA_VISIBLE_DEVICES=0 python main.py --victim RN101 --output-dim 512

svhn_at.py

@@ -0,0 +1,182 @@
+from PIL import Image
+import numpy as np
+import timm
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import torchvision.transforms as transforms
+from torch.utils.data import Dataset, DataLoader
+from torch.optim.lr_scheduler import MultiStepLR, StepLR
+
+from art.estimators.classification import PyTorchClassifier
+from art.data_generators import PyTorchDataGenerator
+from art.defences.trainer import AdversarialTrainer
+from art.attacks.evasion import ProjectedGradientDescent
+from datasets import load_dataset
+from torchvision.transforms import (CenterCrop,
+                                    Compose,
+                                    Normalize,
+                                    RandomHorizontalFlip,
+                                    RandomResizedCrop,
+                                    Resize,
+                                    ToTensor)
+from tensorflow.keras.utils import to_categorical
+from transformers import ViTImageProcessor
+
+processor = ViTImageProcessor.from_pretrained("google/vit-base-patch16-224-in21k")
+IMAGENET_DEFAULT_MEAN = processor.image_mean
+IMAGENET_DEFAULT_STD = processor.image_std
+
+size = processor.size["height"]
+
+
+model = timm.create_model("timm/vit_base_patch16_224.orig_in21k_ft_in1k",
+pretrained=False)
+model.head = nn.Linear(model.head.in_features, 10)
+model.load_state_dict(
+    torch.hub.load_state_dict_from_url(
+        "https://huggingface.co/edadaltocg/vit_base_patch16_224_in21k_ft_svhn/resolve/main/pytorch_model.bin",
+        map_location="cuda",
+        file_name="vit_base_patch16_224_in21k_ft_svhn.pth",
+    )
+)
+
+train_ds = load_dataset('ufldl-stanford/svhn', "cropped_digits", split="train")
+test_ds = load_dataset('ufldl-stanford/svhn', "cropped_digits", split="test")
+splits = train_ds.train_test_split(test_size=0.1)
+train_ds = splits['train']
+val_ds = splits['test']
+
+train_size=len(train_ds)
+test_size=len(test_ds)
+
+normalize = Normalize(mean=IMAGENET_DEFAULT_MEAN, std=IMAGENET_DEFAULT_STD)
+_train_transforms = Compose(
+        [
+            RandomResizedCrop(size),
+            RandomHorizontalFlip(),
+            ToTensor(),
+            normalize,
+        ]
+    )
+
+_val_transforms = Compose(
+        [
+            Resize(size),
+            CenterCrop(size),
+            ToTensor(),
+            normalize,
+        ]
+    )
+def train_transforms(examples):
+    examples['pixel_values'] = [_train_transforms(image.convert("RGB")) for image in examples['image']]
+    return examples
+
+def val_transforms(examples):
+    examples['pixel_values'] = [_val_transforms(image.convert("RGB")) for image in examples['image']]
+    return examples
+
+train_ds.set_transform(train_transforms)
+val_ds.set_transform(val_transforms)
+test_ds.set_transform(val_transforms)
+
+def collate_fn(examples):
+    pixel_values = torch.stack([example["pixel_values"] for example in examples])
+    labels = torch.tensor([example["label"] for example in examples])
+    return pixel_values,labels
+
+train_batch_size = 32
+eval_batch_size = 32
+
+def dataset2np(dataset):
+    X = []
+    Y = []
+    for i in range(int(2000)):
+        x,y = dataset[i]["pixel_values"], dataset[i]["label"]
+        y=to_categorical(y,num_classes=10)
+        X.append(x.detach().numpy())
+        Y.append(y)
+    X = np.array(X).astype("float32")
+    Y = np.array(Y).astype("float32")
+    return X,Y
+
+train_dataloader = DataLoader(train_ds, shuffle=True, collate_fn=collate_fn, batch_size=train_batch_size)
+val_dataloader = DataLoader(val_ds, collate_fn=collate_fn, batch_size=eval_batch_size)
+test_dataloader = DataLoader(test_ds, collate_fn=collate_fn, batch_size=eval_batch_size)
+x_test, y_test=dataset2np(test_ds)
+
+
+opt = torch.optim.SGD(model.parameters(), lr=0.1, momentum=0.9, weight_decay=2e-4)
+lr_scheduler = StepLR(opt, step_size=3, gamma=0.1)
+
+criterion = nn.CrossEntropyLoss()
+
+# Step 3: Create the ART classifier
+
+classifier = PyTorchClassifier(
+    model=model,
+    clip_values=(0.0, 1.0),
+    loss=criterion,
+    optimizer=opt,
+    input_shape=(3, size, size),
+    nb_classes=10,
+)
+
+attack = ProjectedGradientDescent(
+    classifier,
+    norm=np.inf,
+    eps=8.0 / 255.0,
+    eps_step=2.0 / 255.0,
+    max_iter=10,
+    targeted=False,
+    num_random_init=1,
+    batch_size=128,
+    verbose=False,
+)
+
+x_test_clean_pred=np.argmax(classifier.predict(x_test), axis=1)
+print(
+    "Accuracy on clean samples before adversarial training: %.2f%%"
+    % (np.sum(x_test_clean_pred == np.argmax(y_test, axis=1)) / x_test.shape[0] * 100)
+)
+
+
+# Step 4: Create the trainer object - AdversarialTrainerTRADESPyTorch
+trainer = AdversarialTrainer(
+    classifier, attack
+)
+
+# Build a Keras image augmentation object and wrap it in ART
+art_datagen = PyTorchDataGenerator(iterator=train_dataloader, size=train_size, batch_size=128)
+
+# Step 5: fit the trainer
+trainer.fit_generator(art_datagen, nb_epochs=50)
+
+
+x_test_pred = np.argmax(classifier.predict(x_test), axis=1)
+print(
+    "Accuracy on benign test samples after adversarial training: %.2f%%"
+    % (np.sum(x_test_pred == np.argmax(y_test, axis=1)) / x_test.shape[0] * 100)
+)
+
+attack_test = ProjectedGradientDescent(
+    classifier,
+    norm=np.inf,
+    eps=8.0 / 255.0,
+    eps_step=2.0 / 255.0,
+    max_iter=20,
+    targeted=False,
+    num_random_init=1,
+    batch_size=128,
+    verbose=False,
+)
+x_test_attack = attack_test.generate(x_test, y=y_test)
+x_test_attack_pred = np.argmax(classifier.predict(x_test_attack), axis=1)
+print(
+    "Accuracy on original PGD adversarial samples after adversarial training: %.2f%%"
+    % (np.sum(x_test_attack_pred == np.argmax(y_test, axis=1)) / x_test.shape[0] * 100)
+)
+torch.save(trainer.classifier.model.state_dict(), 'svhn-pgd.pth')
+print(
+    "Save the AT model! "
+)

train/hash_train.py

@@ -116,8 +116,6 @@ class Trainer(TrainBase):
                    beta=10 ,epsilon=0.03125, alpha=3/255, num_iter=1500, temperature=0.05):
 
         delta = torch.zeros_like(image,requires_grad=True)
-        # one=torch.zeros_like(positive)
-        # alienation_loss = nn.TripletMarginLoss(margin=1.0, p=2, eps=1e-7)
         for i in range(num_iter):
             self.model.zero_grad()
             anchor=self.model.encode_image(image+delta)
@@ -185,18 +183,18 @@ class Trainer(TrainBase):
 
         mAP_t=cal_map(adv_img,adv_labels,retrieval_txt,retrieval_labels)
         # pr=cal_pr(retrieval_txt,adv_img,query_labels,retrieval_labels)
-        # pr_t=cal_pr(retrieval_txt,adv_img,adv_labels,retrieval_labels)
+        pr_t=cal_pr(retrieval_txt,adv_img,retrieval_labels,adv_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
+            'r_l': retrieval_labels,
             # 'q_l':query_labels
             # 'pr': pr,
-            # 'pr_t': pr_t
+            'pr_t': pr_t
         }
-        scio.savemat(os.path.join(save_dir, str(self.args.output_dim) + "-adv-" + self.args.dataset  + ".mat"), result_dict)
+        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!")
         
         
@@ -269,8 +267,8 @@ class Trainer(TrainBase):
         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)
+        pr_i2t=cal_pr(retrieval_txt,query_img,retrieval_labels,query_labels)
+        pr_t2i=cal_pr(retrieval_img,query_txt,retrieval_labels,query_labels)
         self.max_mapt2i = max(self.max_mapt2i, mAPi2t)
         self.logger.info(f">>>>>> MAP(i->t): {mAPi2t}, MAP(t->i): {mAPt2i}")
         result_dict = {
@@ -279,35 +277,14 @@ class Trainer(TrainBase):
             'r_img': retrieval_img,
             'r_txt': retrieval_txt,
             'q_l': query_labels,
-            'r_l': retrieval_labels
-            # 'pr_i2t': pr_i2t,
-            # 'pr_t2i': pr_t2i
+            'r_l': retrieval_labels,
+            'pr_i2t': pr_i2t,
+            'pr_t2i': pr_t2i
         }
-        scio.savemat(os.path.join(save_dir, str(self.args.output_dim) + "-ours-" + self.args.dataset  + ".mat"), result_dict)
+        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 valid(self, epoch):
-    #     self.logger.info("Valid.")
-    #     self.change_state(mode="valid")
-    #     query_img, query_txt = self.get_code(self.query_loader, self.args.query_num) if self.args.hash_layer == "select" else super().get_code(self.query_loader, self.args.query_num)
-    #     retrieval_img, retrieval_txt = self.get_code(self.retrieval_loader, self.args.retrieval_num) if self.args.hash_layer == "select" else super().get_code(self.retrieval_loader, self.args.retrieval_num)
-    #     # print("get all code")
-    #     mAPi2t = calc_map_k(query_img, retrieval_txt, self.query_labels, self.retrieval_labels, None, self.rank)
-    #     # print("map map")
-    #     mAPt2i = calc_map_k(query_txt, retrieval_img, self.query_labels, self.retrieval_labels, None, self.rank)
-    #     mAPi2i = calc_map_k(query_img, retrieval_img, self.query_labels, self.retrieval_labels, None, self.rank)
-    #     mAPt2t = calc_map_k(query_txt, retrieval_txt, self.query_labels, self.retrieval_labels, None, self.rank)
-    #     if self.max_mapi2t < mAPi2t:
-    #         self.best_epoch_i = epoch
-    #         self.save_mat(query_img, query_txt, retrieval_img, retrieval_txt, mode_name="i2t")
-    #     self.max_mapi2t = max(self.max_mapi2t, mAPi2t)
-    #     if self.max_mapt2i < mAPt2i:
-    #         self.best_epoch_t = epoch
-    #         self.save_mat(query_img, query_txt, retrieval_img, retrieval_txt, mode_name="t2i")
-    #     self.max_mapt2i = max(self.max_mapt2i, mAPt2i)
-    #     self.logger.info(f">>>>>> [{epoch}/{self.args.epochs}], MAP(i->t): {mAPi2t}, MAP(t->i): {mAPt2i}, MAP(t->t): {mAPt2t}, MAP(i->i): {mAPi2i}, \
-    #                 MAX MAP(i->t): {self.max_mapi2t}, MAX MAP(t->i): {self.max_mapt2i}")
 
     def save_mat(self, query_img, query_txt, retrieval_img, retrieval_txt, mode_name="i2t"):
 

train/text_train.py

@@ -1,3 +1,4 @@
+import copy
 from torch.nn.modules import loss
 # from model.hash_model import DCMHT as DCMHT
 import os
@@ -11,38 +12,107 @@ 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 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
-# from transformers import BertModel
+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")
 
-def clamp(delta, clean_imgs):
+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)
 
-    clamp_imgs = (delta.data + clean_imgs.data).clamp(0, 1)
-    clamp_delta = clamp_imgs - clean_imgs.data
 
-    return clamp_delta
+def text_filter(text):
+    text = re.findall(r"<|startoftext|>(.+)<|endoftext|>", text)
+    text = text[2]
+    text = re.sub(r'</w>', ' ', 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):
+                 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
+        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 = model_clip
         self.model.eval()
         self.model.float()
 
@@ -52,14 +122,14 @@ class Trainer(TrainBase):
         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)
+        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()
@@ -67,125 +137,313 @@ class Trainer(TrainBase):
         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
-            )
+            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
-            )
+            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
-        
-    
+            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=[]
+        image_train = []
+        label_train = []
         for image, text, label, index in self.train_loader:
-            image=image.to(device, non_blocking=True)
+            # 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)
+            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_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,  image, negetive_code,negetive_mean,negative_var, positive_code,positive_mean,positive_var,
-                   beta=10 ,epsilon=0.03125, alpha=3/255, num_iter=1500, temperature=0.05):
+            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
 
-        delta = torch.zeros_like(image,requires_grad=True)
-        # one=torch.zeros_like(positive)
-        # alienation_loss = nn.TripletMarginLoss(margin=1.0, p=2, eps=1e-7)
-        for i in range(num_iter):
-            self.model.zero_grad()
-            anchor=self.model.encode_image(image+delta)
-            loss1=F.triplet_margin_with_distance_loss(anchor, positive_code,negetive_code, distance_function=nn.CosineSimilarity())
-            negative_dist=(anchor-negetive_mean)**2 / negative_var
-            positive_dist=(anchor-positive_mean)**2 /positive_var
-            negatives=torch.exp(negative_dist / temperature) 
-            positives= torch.exp(positive_dist / temperature)
-            loss=  torch.log(positives/(positives+negatives)).mean() + beta* loss1
-            loss.backward(retain_graph=True)
-            delta.data = delta - alpha * delta.grad.detach().sign()
-            delta.data =clamp(delta, image).clamp(-epsilon, epsilon)
-            delta.grad.zero_()
-        adv_code=self.model.encode_image(image+delta)
-        return delta.detach() , adv_code
-    
     def train_epoch(self):
-        self.change_state(mode="valid")
+        # 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=[]
+        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)
-            
+            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)
-            
-        
-            delta, adv_code=self.target_adv(image,negetive_code,negetive_mean,negative_var,
-                                  positive_code,positive_mean,positive_var)
+            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)
+        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_img, 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)
-        # pr=cal_pr(retrieval_txt,adv_img,query_labels,retrieval_labels)
-        # pr_t=cal_pr(retrieval_txt,adv_img,adv_labels,retrieval_labels)
+        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,
@@ -196,13 +454,9 @@ class Trainer(TrainBase):
             # 'pr': pr,
             # 'pr_t': pr_t
         }
-        scio.savemat(os.path.join(save_dir, str(self.args.output_dim) + "-adv-" + self.args.dataset  + ".mat"), result_dict)
+        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.")
@@ -212,9 +466,8 @@ class Trainer(TrainBase):
             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}")
-
-    
+        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:
 
@@ -241,25 +494,19 @@ class Trainer(TrainBase):
                 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):
+
+        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) 
+        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()
@@ -267,8 +514,8 @@ class Trainer(TrainBase):
         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)
+        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)
@@ -281,31 +528,11 @@ class Trainer(TrainBase):
             'q_l': query_labels,
             'r_l': retrieval_labels
         }
-        scio.savemat(os.path.join(save_dir, str(self.args.output_dim) + "-ours-" + self.args.dataset  + ".mat"), result_dict)
+        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 valid(self, epoch):
-    #     self.logger.info("Valid.")
-    #     self.change_state(mode="valid")
-    #     query_img, query_txt = self.get_code(self.query_loader, self.args.query_num) if self.args.hash_layer == "select" else super().get_code(self.query_loader, self.args.query_num)
-    #     retrieval_img, retrieval_txt = self.get_code(self.retrieval_loader, self.args.retrieval_num) if self.args.hash_layer == "select" else super().get_code(self.retrieval_loader, self.args.retrieval_num)
-    #     # print("get all code")
-    #     mAPi2t = calc_map_k(query_img, retrieval_txt, self.query_labels, self.retrieval_labels, None, self.rank)
-    #     # print("map map")
-    #     mAPt2i = calc_map_k(query_txt, retrieval_img, self.query_labels, self.retrieval_labels, None, self.rank)
-    #     mAPi2i = calc_map_k(query_img, retrieval_img, self.query_labels, self.retrieval_labels, None, self.rank)
-    #     mAPt2t = calc_map_k(query_txt, retrieval_txt, self.query_labels, self.retrieval_labels, None, self.rank)
-    #     if self.max_mapi2t < mAPi2t:
-    #         self.best_epoch_i = epoch
-    #         self.save_mat(query_img, query_txt, retrieval_img, retrieval_txt, mode_name="i2t")
-    #     self.max_mapi2t = max(self.max_mapi2t, mAPi2t)
-    #     if self.max_mapt2i < mAPt2i:
-    #         self.best_epoch_t = epoch
-    #         self.save_mat(query_img, query_txt, retrieval_img, retrieval_txt, mode_name="t2i")
-    #     self.max_mapt2i = max(self.max_mapt2i, mAPt2i)
-    #     self.logger.info(f">>>>>> [{epoch}/{self.args.epochs}], MAP(i->t): {mAPi2t}, MAP(t->i): {mAPt2i}, MAP(t->t): {mAPt2t}, MAP(i->i): {mAPi2i}, \
-    #                 MAX MAP(i->t): {self.max_mapi2t}, MAX MAP(t->i): {self.max_mapt2i}")
 
     def save_mat(self, query_img, query_txt, retrieval_img, retrieval_txt, mode_name="i2t"):
 
@@ -327,7 +554,7 @@ class Trainer(TrainBase):
             'q_l': query_labels,
             'r_l': retrieval_labels
         }
-        scio.savemat(os.path.join(save_dir, str(self.args.output_dim) + "-ours-" + self.args.dataset + "-" + mode_name + ".mat"), result_dict)
+        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!")
-
-

utils/get_args.py

@@ -9,19 +9,22 @@ def get_args():
     parser.add_argument("--save-dir", type=str, default="./result/64-bit")
     parser.add_argument("--clip-path", type=str, default="./ViT-B-32.pt", help="pretrained clip path.")
     parser.add_argument("--pretrained", type=str, default="")
-    parser.add_argument("--dataset", type=str, default="flickr25k", help="choise from [coco, mirflckr25k, nuswide]")
+    parser.add_argument("--dataset", type=str, default="coco", help="choise from [coco, mirflckr25k, nuswide]")
     parser.add_argument("--index-file", type=str, default="index.mat")
     parser.add_argument("--caption-file", type=str, default="caption.mat")
     parser.add_argument("--label-file", type=str, default="label.mat")
     parser.add_argument("--similarity-function", type=str, default="euclidean", help="choise form [cosine, euclidean]")
     parser.add_argument("--loss-type", type=str, default="l2", help="choise form [l1, l2]")
     parser.add_argument('--victim', default='ViT-B/16', choices=['ViT-L/14', 'ViT-B/16', 'ViT-B/32', 'RN50', 'RN101'])
-    # parser.add_argument("--test-caption-file", type=str, default="./data/test/captions.mat")
-    # parser.add_argument("--test-label-file", type=str, default="./data/test/label.mat")
+    parser.add_argument("--text_encoder", type=str, default="bert-base-uncased")
+    parser.add_argument("--topk", type=int, default=10)
+    parser.add_argument("--num-perturbation", type=int, default=3)
     parser.add_argument("--txt-dim", type=int, default=1024)
     parser.add_argument("--output-dim", type=int, default=512)
     parser.add_argument("--epochs", type=int, default=100)
     parser.add_argument("--max-words", type=int, default=77)
+    parser.add_argument("--max-candidate", type=int, default=7)
+    parser.add_argument("--enable-bpe", type=bool, default=False)
     parser.add_argument("--resolution", type=int, default=224)
     parser.add_argument("--batch-size", type=int, default=8)
     parser.add_argument("--num-workers", type=int, default=4)
@@ -30,7 +33,7 @@ def get_args():
     parser.add_argument("--lr-decay-freq", type=int, default=5)
     parser.add_argument("--display-step", type=int, default=50)
     parser.add_argument("--seed", type=int, default=1814)
-
+    parser.add_argument("--attack-thred", type=float, default=0.05)
     parser.add_argument("--lr", type=float, default=0.001)
     parser.add_argument("--lr-decay", type=float, default=0.9)
     parser.add_argument("--clip-lr", type=float, default=0.00001)