增加vgg网络水印验证代码

verify_cifar10_vgg16.py

@@ -5,7 +5,8 @@ from matplotlib import pyplot as plt
 from keras.layers import Conv2D, BatchNormalization, Activation, MaxPool2D, Dropout, Flatten, Dense
 from keras import Model
 
-from class5.CIFAR10_CNN.watermark_regularizers import WatermarkRegularizer
+from tf_watermark.tf_watermark_regularizers import WatermarkRegularizer
+from tf_watermark.tf_watermark_utils import save_wmark_signatures, get_layer_weights_and_predicted
 
 np.set_printoptions(threshold=np.inf)
 
@@ -13,12 +14,17 @@ cifar10 = tf.keras.datasets.cifar10
 (x_train, y_train), (x_test, y_test) = cifar10.load_data()
 x_train, x_test = x_train / 255.0, x_test / 255.0
 
-# 程序参数
-target_blk_id = 0  # 目标层ID
-embed_dim = 256  # 水印长度
+# 初始化参数
 scale = 0.01  # 正则化项偏置系数
-wtype = 'random'  # 水印类型
-randseed = 'none'
+randseed = 5  # 投影矩阵生成随机数种子
+embed_dim = 768  # 密钥长度
+np.random.seed(5)
+b = np.random.randint(low=0, high=2, size=(1, embed_dim))  # 生成模拟随机密钥
+epoch = 25
+
+# 初始化水印正则化器
+watermark_regularizer = WatermarkRegularizer(scale, b)
+
 
 class VGG16(Model):
     def __init__(self):
@@ -44,7 +50,7 @@ class VGG16(Model):
         self.c5 = Conv2D(filters=256, kernel_size=(3, 3), padding='same')
         self.b5 = BatchNormalization()  # BN层1
         self.a5 = Activation('relu')  # 激活层1
-        self.c6 = Conv2D(filters=256, kernel_size=(3, 3), padding='same')
+        self.c6 = Conv2D(filters=256, kernel_size=(3, 3), padding='same', kernel_regularizer=watermark_regularizer)
         self.b6 = BatchNormalization()  # BN层1
         self.a6 = Activation('relu')  # 激活层1
         self.c7 = Conv2D(filters=256, kernel_size=(3, 3), padding='same')
@@ -150,14 +156,6 @@ class VGG16(Model):
 
 model = VGG16()
 
-# 初始化水印过程
-b = np.ones((1, embed_dim))
-wmark_regularizer = WatermarkRegularizer(scale, b, wtype=wtype, randseed=randseed)
-
-# 添加正则化项到指定层的权重
-selected_layer = model.get_layer(index=target_blk_id)
-selected_layer.kernel_regularizer = wmark_regularizer
-
 model.compile(optimizer='adam',
               loss=tf.keras.losses.SparseCategoricalCrossentropy(from_logits=False),
               metrics=['sparse_categorical_accuracy'])
@@ -171,17 +169,20 @@ cp_callback = tf.keras.callbacks.ModelCheckpoint(filepath=checkpoint_save_path,
                                                  save_weights_only=True,
                                                  save_best_only=True)
 
-history = model.fit(x_train, y_train, batch_size=32, epochs=20, validation_data=(x_test, y_test), validation_freq=1,
+history = model.fit(x_train, y_train, batch_size=64, epochs=epoch, validation_data=(x_test, y_test), validation_freq=1,
                     callbacks=[cp_callback])
 model.summary()
 
-# print(model.trainable_variables)
-file = open('./weights.txt', 'w')
-for v in model.trainable_variables:
-    file.write(str(v.name) + '\n')
-    file.write(str(v.shape) + '\n')
-    file.write(str(v.numpy()) + '\n')
-file.close()
+###############################################    verify watermarker ###################################
+# 保存投影矩阵和密钥
+save_wmark_signatures(model)
+target_layer = model.get_layer(index=19)
+layer_weights, pred_bparam = get_layer_weights_and_predicted(target_layer)
+print("b_param:")
+print(b)
+print("pred_bparam:")
+print(pred_bparam)
+print(np.sum(b != pred_bparam))
 
 ###############################################    show   ###############################################