项目初始化

.gitignore

@@ -46,3 +46,4 @@ com_crashlytics_export_strings.xml
 crashlytics.properties
 crashlytics-build.properties
 
+__pycache__

models/alexnet.py

@@ -0,0 +1,56 @@
+import torch
+import torch.nn as nn
+
+
+# 定义修改后的AlexNet模型
+class AlexNet(nn.Module):
+    def __init__(self):
+        super(AlexNet, self).__init__()
+        # 定义每一个就卷积层
+        self.layer1 = nn.Sequential(
+            # 卷积层  #输入图像为1*28*28
+            nn.Conv2d(3, 32, kernel_size=3, padding=1),
+            # 池化层
+            nn.MaxPool2d(kernel_size=2, stride=2),  # 池化层特征图通道数不改变，每个特征图的分辨率变小
+            # 激活函数Relu
+            nn.ReLU(inplace=True),
+        )
+
+        self.layer2 = nn.Sequential(
+            nn.Conv2d(32, 64, kernel_size=3, padding=1),
+            nn.MaxPool2d(kernel_size=2, stride=2),
+            nn.ReLU(inplace=True),
+        )
+
+        self.layer3 = nn.Sequential(
+            nn.Conv2d(64, 128, kernel_size=3, padding=1),
+        )
+
+        self.layer4 = nn.Sequential(
+            nn.Conv2d(128, 256, kernel_size=3, padding=1),
+        )
+        self.layer5 = nn.Sequential(
+            nn.Conv2d(256, 256, kernel_size=3, padding=1),
+            nn.MaxPool2d(kernel_size=3, stride=2),
+            nn.ReLU(inplace=True),
+        )
+
+        # 定义全连接层
+        self.fc1 = nn.Linear(256 * 3 * 3, 1024)
+        self.fc2 = nn.Linear(1024, 512)
+        self.fc3 = nn.Linear(512, 10)
+        # 对应十个类别的输出
+
+    def forward(self, x):
+        x = self.layer1(x)
+        x = self.layer2(x)
+        x = self.layer3(x)
+        x = self.layer4(x)
+        x = self.layer5(x)
+        # print(x.shape)
+        x = x.view(-1, 256 * 3 * 3)
+        x = self.fc1(x)
+        x = self.fc2(x)
+        x = self.fc3(x)
+
+        return x

models/embeder.py

@@ -0,0 +1,17 @@
+import torch.nn as nn
+
+
+class WatermarkEmbeder(nn.Module):
+
+    def __init__(self, pct_dim, wm_dim):
+        super(WatermarkEmbeder, self).__init__()
+
+        self.model = nn.Sequential(
+            nn.Linear(pct_dim, 256, bias=False),
+            nn.ReLU(inplace=True),
+            nn.Linear(256, wm_dim, bias=False),
+            nn.Sigmoid()
+        )
+
+    def forward(self, x):
+        return self.model(x)

models/googlenet.py

@@ -0,0 +1,82 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+
+# 定义GoogLeNet模型
+class InceptionModule(nn.Module):
+    def __init__(self, in_channels, out1, out2_in, out2, out3_in, out3, out4):
+        super(InceptionModule, self).__init__()
+        self.branch1 = nn.Sequential(
+            nn.Conv2d(in_channels, out1, kernel_size=1),
+        )
+        self.branch2 = nn.Sequential(
+            nn.Conv2d(in_channels, out2_in, kernel_size=1),
+            nn.Conv2d(out2_in, out2, kernel_size=3, padding=1),
+        )
+        self.branch3 = nn.Sequential(
+            nn.Conv2d(in_channels, out3_in, kernel_size=1),
+            nn.Conv2d(out3_in, out3, kernel_size=5, padding=2),
+        )
+        self.branch4 = nn.Sequential(
+            nn.MaxPool2d(kernel_size=3, stride=1, padding=1),
+            nn.Conv2d(in_channels, out4, kernel_size=1),
+        )
+
+    def forward(self, x):
+        out1 = self.branch1(x)
+        out2 = self.branch2(x)
+        out3 = self.branch3(x)
+        out4 = self.branch4(x)
+        return torch.cat([out1, out2, out3, out4], 1)
+
+
+class GoogLeNet(nn.Module):
+    def __init__(self, num_classes=10):
+        super(GoogLeNet, self).__init__()
+        self.conv1 = nn.Sequential(
+            nn.Conv2d(3, 64, kernel_size=7, stride=2, padding=3),
+            nn.ReLU(inplace=True),
+            nn.MaxPool2d(kernel_size=3, stride=2, padding=1),
+        )
+        self.conv2 = nn.Sequential(
+            nn.Conv2d(64, 64, kernel_size=1),
+            nn.ReLU(inplace=True),
+            nn.Conv2d(64, 192, kernel_size=3, padding=1),
+            nn.ReLU(inplace=True),
+            nn.MaxPool2d(kernel_size=3, stride=2, padding=1),
+        )
+        self.inception3a = InceptionModule(192, 64, 96, 128, 16, 32, 32)
+        self.inception3b = InceptionModule(256, 128, 128, 192, 32, 96, 64)
+        self.maxpool = nn.MaxPool2d(kernel_size=3, stride=2, padding=1)
+        self.inception4a = InceptionModule(480, 192, 96, 208, 16, 48, 64)
+        self.inception4b = InceptionModule(512, 160, 112, 224, 24, 64, 64)
+        self.inception4c = InceptionModule(512, 128, 128, 256, 24, 64, 64)
+        self.inception4d = InceptionModule(512, 112, 144, 288, 32, 64, 64)
+        self.inception4e = InceptionModule(528, 256, 160, 320, 32, 128, 128)
+        self.maxpool = nn.MaxPool2d(kernel_size=3, stride=2, padding=1)
+        self.inception5a = InceptionModule(832, 256, 160, 320, 32, 128, 128)
+        self.inception5b = InceptionModule(832, 384, 192, 384, 48, 128, 128)
+        self.avgpool = nn.AdaptiveAvgPool2d((1, 1))
+        self.dropout = nn.Dropout(0.4)
+        self.fc = nn.Linear(1024, num_classes)
+
+    def forward(self, x):
+        x = self.conv1(x)
+        x = self.conv2(x)
+        x = self.inception3a(x)
+        x = self.inception3b(x)
+        x = self.maxpool(x)
+        x = self.inception4a(x)
+        x = self.inception4b(x)
+        x = self.inception4c(x)
+        x = self.inception4d(x)
+        x = self.inception4e(x)
+        x = self.maxpool(x)
+        x = self.inception5a(x)
+        x = self.inception5b(x)
+        x = self.avgpool(x)
+        x = x.view(x.size(0), -1)
+        x = self.dropout(x)
+        x = self.fc(x)
+        return x

models/lenet.py

@@ -0,0 +1,26 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+
+class LeNet(nn.Module):
+    def __init__(self):
+        super(LeNet, self).__init__()
+        self.conv1 = nn.Conv2d(3, 16, 5)
+        self.pool1 = nn.MaxPool2d(2, 2)
+        self.conv2 = nn.Conv2d(16, 32, 5)
+        self.pool2 = nn.MaxPool2d(2, 2)
+        self.fc1 = nn.Linear(32 * 5 * 5, 120)
+        self.fc2 = nn.Linear(120, 84)
+        self.fc3 = nn.Linear(84, 10)
+
+    def forward(self, x):
+        x = F.relu(self.conv1(x))  # input(3,32,32) output(16,28,28)
+        x = self.pool1(x)  # output(16，14，14)
+        x = F.relu(self.conv2(x))  # output(32,10.10)
+        x = self.pool2(x)  # output(32,5,5)
+        x = x.view(-1, 32 * 5 * 5)  # output(5*5*32)
+        x = F.relu(self.fc1(x))  # output(120)
+        x = F.relu(self.fc2(x))  # output(84)
+        x = self.fc3(x)  # output(10)
+        return x

models/vgg16.py

@@ -0,0 +1,36 @@
+import torch.nn as nn
+
+
+# 定义VGG16的网络结构
+class VGGNet(nn.Module):
+    def __init__(self, num_classes=10):
+        super(VGGNet, self).__init__()
+        self.features = self._make_layers()
+        self.classifier = nn.Sequential(
+            nn.Linear(512 * 1 * 1, 512),
+            nn.ReLU(inplace=True),
+            nn.Dropout(),
+            nn.Linear(512, 512),
+            nn.ReLU(inplace=True),
+            nn.Dropout(),
+            nn.Linear(512, num_classes)
+        )
+
+    def _make_layers(self):
+        layers = []
+        in_channels = 3
+        cfg = [64, 64, 'M', 128, 128, 'M', 256, 256, 256, 'M', 512, 512, 512, 'M', 512, 512, 512, 'M']
+        for v in cfg:
+            if v == 'M':
+                layers += [nn.MaxPool2d(kernel_size=2, stride=2)]
+            else:
+                conv2d = nn.Conv2d(in_channels, v, kernel_size=3, padding=1)
+                layers += [conv2d, nn.ReLU(inplace=True)]
+                in_channels = v
+        return nn.Sequential(*layers)
+
+    def forward(self, x):
+        x = self.features(x)
+        x = x.view(x.size(0), -1)
+        x = self.classifier(x)
+        return x

run.py

@@ -0,0 +1,45 @@
+import numpy as np
+import onnx
+import onnxruntime
+import torchvision
+import torch
+from torchvision.transforms import transforms
+from onnx import numpy_helper
+
+transform_test = transforms.Compose([
+    transforms.ToTensor(),
+    transforms.Normalize((0.485, 0.456, 0.406), (0.229, 0.224, 0.225)),
+])
+
+testset = torchvision.datasets.CIFAR10(root='./data', train=False, download=True, transform=transform_test)
+testloader = torch.utils.data.DataLoader(testset, batch_size=1, shuffle=False)
+
+ort_session = onnxruntime.InferenceSession('host_dnn.onnx')
+ort_session_embeder = onnxruntime.InferenceSession('embeder.onnx')
+# 获取目标模型权重
+model = onnx.load('host_dnn.onnx')
+weights = model.graph.initializer
+weight = numpy_helper.to_array(weights[0])
+weight = weight.reshape((1, -1))
+
+correct = 0
+total = 0
+# 读取随机密钥
+b = np.load('b.npy')
+
+ort_input_embeder = {'input': weight}
+ort_output_embeder = ort_session_embeder.run(['output'], ort_input_embeder)
+
+result = (ort_output_embeder == b).all()
+
+for data in testloader:
+    inputs, labels = data
+    inputs, labels = inputs.numpy(), labels.numpy()
+    ort_inputs = {'input': inputs}
+    ort_outputs = ort_session.run(['output'], ort_inputs)
+    result = labels == np.argmax(ort_outputs)
+    if result:
+        correct = correct + 1
+    total = total + 1
+
+print(f'acc={correct / total}')

train.py

@@ -0,0 +1,191 @@
+import numpy as np
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import torchvision
+import torchvision.transforms as transforms
+from matplotlib import pyplot as plt
+from torch import optim
+from tqdm import tqdm  # 导入tqdm
+
+from models.alexnet import AlexNet
+from models.embeder import WatermarkEmbeder
+from models.googlenet import GoogLeNet
+from models.lenet import LeNet
+from models.vgg16 import VGGNet
+
+# 参数
+batch_size = 500
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+num_epochs = 40
+wm_length = 1024
+
+# 设置随机数种子
+np.random.seed(1)
+lambda1 = 0.05
+b = np.random.randint(low=0, high=2, size=(1, wm_length))  # 生成模拟随机密钥
+np.save('b.npy', b)
+b = nn.Parameter(torch.tensor(b, dtype=torch.float32).to(device), requires_grad=False)
+b.requires_grad = False
+
+# 数据预处理和加载
+transform_train = transforms.Compose([
+    transforms.RandomCrop(32, padding=4),
+    transforms.RandomHorizontalFlip(),
+    transforms.ToTensor(),
+    transforms.Normalize((0.485, 0.456, 0.406), (0.229, 0.224, 0.225)),
+])
+
+transform_test = transforms.Compose([
+    transforms.ToTensor(),
+    transforms.Normalize((0.485, 0.456, 0.406), (0.229, 0.224, 0.225)),
+])
+
+trainset = torchvision.datasets.CIFAR10(root='./data', train=True, download=True, transform=transform_train)
+trainloader = torch.utils.data.DataLoader(trainset, batch_size=batch_size, shuffle=True)
+
+testset = torchvision.datasets.CIFAR10(root='./data', train=False, download=True, transform=transform_test)
+testloader = torch.utils.data.DataLoader(testset, batch_size=100, shuffle=False)
+
+# 创建VGGNet模型实例
+# model = VGGNet(num_classes=10).to(device)
+# model = GoogLeNet(num_classes=10).to(device)
+# model = LeNet().to(device)
+model = AlexNet().to(device)
+
+# 打印模型结构
+print(model)
+
+# 断点续训
+# model.load_state_dict(torch.load("./result/host_dnn.pt"))
+
+# 获取指定层权重信息
+# vgg16
+# weight = model.features[0].weight.view(1, -1)
+# googleNet
+# weight = model.conv1[0].weight.view(1, -1)
+# leNet
+# weight = model.conv1.weight.view(1, -1)
+# AlexNet
+weight = model.layer1[0].weight.view(1, -1)
+
+# 获取权重向量长度
+pct_dim = np.prod(weight.shape[1:4])
+# 创建水印嵌入模型实例
+model_embeder = WatermarkEmbeder(pct_dim, wm_length).to(device)
+# model_embeder.load_state_dict(torch.load("./result/embeder.pt"))
+
+# 定义目标模型损失函数和优化器
+criterion = nn.CrossEntropyLoss()
+
+# 目标模型使用Adam优化器
+optimizer = optim.Adam(model.parameters(), lr=1e-4)  # 调整学习率
+optimizer_embeder = optim.Adam(model_embeder.parameters(), lr=0.5)
+
+# 初始化空列表以存储准确度和损失
+train_accs = []
+train_losses = []
+torch.autograd.set_detect_anomaly(True)
+
+for epoch in range(num_epochs):
+    model_embeder.train()
+    model.train()
+    running_loss = 0.0
+    correct = 0
+    total = 0
+
+    # 使用tqdm创建进度条
+    with (tqdm(total=len(trainloader), desc=f"Epoch {epoch + 1}", unit="batch") as pbar):
+        for i, data in enumerate(trainloader, 0):
+            inputs, labels = data
+            inputs, labels = inputs.to(device), labels.to(device)
+
+            optimizer_embeder.zero_grad()
+            outputs_embeder = model_embeder(weight)
+            loss_embeder = F.binary_cross_entropy(outputs_embeder, b)
+            # loss_embeder.backward(retain_graph=True)
+            # optimizer_embeder.step()
+
+            optimizer.zero_grad()
+            outputs = model(inputs)
+            loss_c = criterion(outputs, labels)
+            loss_h = loss_c + lambda1 * loss_embeder
+            loss_h.backward()
+            optimizer.step()
+            optimizer_embeder.step()
+
+            running_loss += loss_h.item()
+
+            _, predicted = torch.max(outputs.data, 1)
+            total += labels.size(0)
+            correct += (predicted == labels).sum().item()
+
+            # 更新进度条
+            pbar.set_postfix(loss=running_loss / (i + 1), loss_em=loss_embeder.item(), acc=100 * correct / total)
+            pbar.update()
+
+        # 计算准确度和损失
+        epoch_acc = 100 * correct / total
+        epoch_loss = running_loss / len(trainloader)
+
+        # 记录准确度和损失值
+        train_accs.append(epoch_acc)
+        train_losses.append(epoch_loss)
+        print(f"Epoch {epoch + 1}, Loss: {epoch_loss}, Accuracy: {epoch_acc}%")
+
+        torch.save(model.state_dict(), "./result/host_dnn.pt")
+        torch.save(model_embeder.state_dict(), "./result/embeder.pt")
+
+        # 导出onnx格式
+        with torch.no_grad():
+            x = torch.randn(1, 3, 32, 32).to(device)
+            torch.onnx.export(model, x, 'host_dnn.onnx', opset_version=11, input_names=['input'],
+                              output_names=['output'])
+            torch.onnx.export(model_embeder, weight, 'embeder.onnx', opset_version=11, input_names=['input'],
+                              output_names=['output'])
+        # 更新学习率
+        # scheduler.step()
+
+    # 测试模型
+    if epoch % 5 == 4:
+        model.eval()
+        correct = 0
+        total = 0
+        with torch.no_grad():
+            for data in testloader:
+                inputs, labels = data
+                inputs, labels = inputs.to(device), labels.to(device)
+                outputs = model(inputs)
+                _, predicted = torch.max(outputs.data, 1)
+                total += labels.size(0)
+                correct += (predicted == labels).sum().item()
+
+        print(f"Accuracy on test set: {(100 * correct / total):.2f}%")
+    # scheduler.step()
+
+print("Finished Training")
+
+# 绘制准确度和损失曲线
+plt.figure(figsize=(12, 4))
+plt.subplot(1, 2, 1)
+plt.plot(train_accs)
+plt.title('Training Accuracy')
+plt.xlabel('Epoch')
+plt.ylabel('Accuracy (%)')
+
+plt.subplot(1, 2, 2)
+plt.plot(train_losses)
+plt.title('Training Loss')
+plt.xlabel('Epoch')
+plt.ylabel('Loss')
+
+plt.tight_layout()
+plt.show()
+
+print("Finished drawing")
+# 测试水印嵌入
+model_embeder.eval()
+outputs_embeder = model_embeder(weight)
+outputs_embeder = (outputs_embeder > 0.5).int()
+wrr_bits = (outputs_embeder != b).sum().item()
+print(f'wrr_bits={wrr_bits}')