initial push

all_AIW/AIWImp/AIMs_files.py

+
+import UCHIDA
+import time
+from utils import *
+from Attacks import *
+
+#####Changement de niveau AIMs doit etre quelque chose qui possède des ports
+class AIM():
+    model=None
+    parameters=None
+    testingDataset=None
+    trainingDataset=None
+
+    output_0=None
+
+    def trainAIM(self, parameters, trainingDataset):
+        model = tv.models.vgg16()
+        model.classifier = nn.Linear(25088, 10)
+        model.load_state_dict(parameters["model_state_dict"])
+        batch_size = 128
+        while batch_size > 0:
+            try:
+                trainloader = torch.utils.data.DataLoader(trainingDataset, batch_size=batch_size, shuffle=False,
+                                                         num_workers=2)
+            except:
+                batch_size = int(batch_size / 2)
+        criterion = nn.CrossEntropyLoss()
+        learning_rate, momentum, weight_decay = 0.01, .9, 5e-4
+        optimizer = optim.SGD([
+            {'params': model.parameters()}
+        ], lr=learning_rate, momentum=momentum, weight_decay=weight_decay)
+        model.train()
+        model.to(device)
+        for epoch in range(10):
+            for i, data in enumerate(trainloader, 0):
+                # split data into the image and its label
+                inputs, labels = data
+                inputs = inputs.to(device)
+                labels = labels.to(device)
+                if inputs.size()[1] == 1:
+                    inputs.squeeze_(1)
+                    inputs = torch.stack([inputs, inputs, inputs], 1)
+                # initialise the optimiser
+                optimizer.zero_grad()
+
+                # forward
+                outputs = model(inputs)
+                # backward
+                loss= criterion(outputs, labels)
+                # watermark
+                loss.backward()
+                # update the optimizer
+                optimizer.step()
+        return model
+
+    def funcAIM(self, model, parameters, testingDataset):
+        '''
+        return the performance of the model in the testing dataset
+        :param parameters: parameter of the network
+        :param testingdataset: testing data
+        :return:
+        '''
+        model.load_state_dict(parameters["model_state_dict"])
+        batch_size=128
+        while batch_size>0:
+            try :
+                testloader=torch.utils.data.DataLoader(testingDataset, batch_size=batch_size, shuffle=False,num_workers=2)
+                return self.testingAIM(model,testloader)
+            except:
+                batch_size=int(batch_size/2)
+
+    def testingAIM(self,AIM,testloader):
+        correct = 0
+        total = 0
+        AIM.to(device)
+        AIM.eval()
+        # torch.no_grad do not train the network
+        with torch.no_grad():
+            for data in testloader:
+                inputs, labels = data
+                if inputs.size()[1] == 1:
+                    inputs.squeeze_(1)
+                    inputs = torch.stack([inputs, inputs, inputs], 1)
+                inputs = inputs.to(device)
+                labels = labels.to(device)
+                outputs = AIM(inputs)
+                if len(outputs) == 2: outputs, _ = outputs
+                _, predicted = torch.max(outputs, 1)
+                total += labels.size(0)
+                correct += (predicted == labels).sum()
+        return 100 - (100 * float(correct) / total)
+
+    def run(self):
+        self.output_0=self.funcAIM(self.model, self.parameters, self.testingDataset)
+
+class WatermarkEmbedder():
+    AIM=None
+    ##
+    output_0 = None
+
+    def funcWatermarkEmbedder(self,model):
+        ### differentiate this function based on trainloader none or not
+        model.to(device)
+        #### to be adapted by
+        tools = UCHIDA.Uchi_tools()
+        weight_name = 'features.19.weight'
+        T = 64
+        watermark = torch.tensor(np.random.choice([0, 1], size=(T), p=[1. / 3, 2. / 3]), device=device)
+        watermarking_dict = {'lambd': 0.1, 'weight_name': weight_name, 'watermark': watermark, "types": 1}
+        if watermarking_dict["types"] == 1:
+            trainset, testset, inference_transform = CIFAR10_dataset()
+            # hyperparameter of training
+            num_epochs = 25
+            batch_size = 128
+            tools.init(model,watermarking_dict)
+            trainloader, testloader = dataloader(trainset, testset, batch_size)
+            criterion = nn.CrossEntropyLoss()
+            learning_rate, momentum, weight_decay = 0.01, .9, 5e-4
+            optimizer = optim.SGD([
+                {'params': model.parameters()}
+            ], lr=learning_rate, momentum=momentum, weight_decay=weight_decay)
+            model.train()
+            epoch = 0
+            print("Launching injection.....")
+            while epoch < num_epochs:
+                print('doing epoch', str(epoch + 1), ".....")
+                loss, loss_nn, loss_w = tools.Embedder_one_step(model, trainloader, optimizer, criterion,
+                                                                watermarking_dict)
+
+                loss = (loss * batch_size / len(trainloader.dataset))
+                loss_nn = (loss_nn * batch_size / len(trainloader.dataset))
+                loss_w = (loss_w * batch_size / len(trainloader.dataset))
+                print(' loss  : %.5f   - loss_wm: %.5f, loss_nn: %.5f  ' % (loss, loss_w, loss_nn))
+
+                epoch += 1
+        else:
+            tools.init(model,watermarking_dict)
+            model = tools.Embedder(model, watermarking_dict)
+
+        ## save the chekcpoints
+        np.save('watermarking_dict.npy', watermarking_dict)
+        torch.save({
+            'model_state_dict': model.state_dict(),
+        }, 'weights')
+        ## load (checkpoint)
+        checkpoints=torch.load('weights', map_location=torch.device('cpu'))
+        return checkpoints
+
+    def run(self):
+        self.output_0=self.funcWatermarkEmbedder(self.AIM)
+
+class Comparator():
+    unwatermarked=None
+    watermarked=None
+    testingDataset = None
+    ##
+    output_0=None
+
+    def funcComparator(self,unwatermarked,watermarked,testingDataset):
+        undertestAIM=AIM()
+
+        res_wm=undertestAIM.funcAIM(watermarked,testingDataset)
+        res_unwm=undertestAIM.funcAIM(unwatermarked,testingDataset)
+        return np.abs(res_unwm-res_wm)/res_wm, res_wm
+
+    def run(self):
+        self.output_0=self.funcComparator(self.unwatermarked, self.watermarked,self.testingDataset)