修改水印代码

tensorflow/tf_watermark_regularizers.py

@@ -11,15 +11,26 @@ class WatermarkRegularizer(Regularizer):
         self.b = b  # 密钥
         self.x = None  # 投影矩阵
         self.v_x = None  # 投影矩阵设为可训练参数
+        self.randseed = randseed
 
     def __call__(self, weight):
         if self.x is None:
             x_rows = np.prod(weight.shape[0:3])
             x_cols = self.b.shape[1]
+            if self.randseed != 'none':
+                np.random.seed(int(self.randseed))
             self.x = np.random.randn(x_rows, x_cols)
             self.v_x = tf.Variable(self.x, trainable=False, dtype=tf.float32)
-            self.b = tf.constant(self.b, dtype=tf.float32)
         weight_mean = tf.reduce_mean(weight, axis=3)
         w = tf.reshape(weight_mean, (1, tf.reduce_prod(weight_mean.shape)))
-        regularized_loss = self.scale * tf.reduce_sum(binary_crossentropy(tf.sigmoid(tf.matmul(w, self.v_x))), self.b)
+        regularized_loss = self.scale * tf.reduce_sum(
+            binary_crossentropy(tf.sigmoid(tf.matmul(w, self.v_x)), tf.cast(self.b, tf.float32)))
         return regularized_loss
+
+    # 获取投影矩阵
+    def get_matrix(self):
+        return self.x
+
+    # 获取签名值
+    def get_signature(self):
+        return self.b

tf_watermark/tf_watermark_utils.py

@@ -0,0 +1,39 @@
+
+import numpy as np
+
+from tf_watermark.tf_watermark_regularizers import WatermarkRegularizer
+
+RESULT_PATH = './result'
+fname_x = RESULT_PATH + '/x.npy'  # 投影矩阵
+fname_b = RESULT_PATH + '/b.npy'  # 密钥
+
+
+def get_wmark_regularizers(model):
+    ret = []
+
+    for i, layer in enumerate(model.layers):
+        if hasattr(layer, "kernel_regularizer"):
+            regularizer = layer.kernel_regularizer
+            if isinstance(regularizer, WatermarkRegularizer):
+                ret.append((i, regularizer))
+    return ret
+
+
+# 保存水印签名的函数
+def save_wmark_signatures(model):
+    for layer_id, wmark_regularizer in get_wmark_regularizers(model):
+        np.save(fname_x, wmark_regularizer.get_matrix())
+        np.save(fname_b, wmark_regularizer.get_signature())
+
+
+# 从文件中获取
+def get_layer_weights_and_predicted(model, checkpoint_save_path, target_blk_id):
+    x = np.load(fname_x)
+    model.load_weights(checkpoint_save_path)
+    # get signature from model weight and matrix
+    target_layer = model.get_layer(index=target_blk_id)
+    layer_weights = target_layer.get_weights()
+    weight = (np.array(layer_weights[0])).mean(axis=3)
+    pred_bparam = np.dot(weight.reshape(1, weight.size), x)  # dot product
+    pred_bparam = np.where(pred_bparam >= 0, 1, 0)
+    return layer_weights[0], pred_bparam

verify_cifar10_alexnet8.py

@@ -5,12 +5,25 @@ from matplotlib import pyplot as plt
 from keras.layers import Conv2D, BatchNormalization, Activation, MaxPool2D, Dropout, Flatten, Dense
 from keras import Model
 
+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)
 
+# 加载数据集
 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
 
+# 初始化参数
+scale = 0.01  # 正则化项偏置系数
+randseed = 5  # 投影矩阵生成随机数种子
+embed_dim = 768  # 密钥长度
+b = np.random.randint(low=0, high=2, size=(1, embed_dim))  # 生成模拟随机密钥
+
+# 初始化水印正则化器
+watermark_regularizer = WatermarkRegularizer(scale, b, randseed=randseed)
+
 
 class AlexNet8(Model):
     def __init__(self):
@@ -27,10 +40,10 @@ class AlexNet8(Model):
 
         self.c3 = Conv2D(filters=384, kernel_size=(3, 3), padding='same',
                          activation='relu')
-                         
+
         self.c4 = Conv2D(filters=384, kernel_size=(3, 3), padding='same',
-                         activation='relu')
-                         
+                         activation='relu', kernel_regularizer=watermark_regularizer)
+
         self.c5 = Conv2D(filters=256, kernel_size=(3, 3), padding='same',
                          activation='relu')
         self.p3 = MaxPool2D(pool_size=(3, 3), strides=2)
@@ -84,17 +97,30 @@ 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=5, validation_data=(x_test, y_test), validation_freq=1,
+history = model.fit(x_train, y_train, batch_size=128, epochs=30, 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()
+# 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()
+
+# 保存投影矩阵和密钥
+save_wmark_signatures(model)
+
+###############################################    verify watermarker ###################################
+
+layer_weights, pred_bparam = get_layer_weights_and_predicted(model, checkpoint_save_path, 9)
+print("b_param:")
+print(b)
+print("pred_bparam:")
+print(pred_bparam)
+print(b == pred_bparam)
+print(np.sum(b != pred_bparam))
 
 ###############################################    show   ###############################################