修改VGG16白盒水印嵌入卷积层位置

README.md

@@ -48,7 +48,7 @@ model_watermark_generate
         │   ├── classification_pytorch_white_embed.py
         │   ├── faster_rcnn_pytorch_black_embed.py
         │   ├── faster_rcnn_pytorch_white_embed.py
-        │   ├── googlenet_pytorch_white_embed.py
+        │   ├── googlenet_vgg16_pytorch_white_embed.py
         │   ├── ssd_pytorch_black_embed.py
         │   ├── ssd_pytorch_white_embed.py
         │   ├── yolox_pytorch_black_embed.py

watermark_generate/controller/watermark_generate_controller.py

@@ -14,7 +14,7 @@ from watermark_generate import logger
 from watermark_generate.tools import secret_label_func
 from watermark_generate.deals import yolox_pytorch_black_embed, yolox_pytorch_white_embed, \
     faster_rcnn_pytorch_black_embed, ssd_pytorch_black_embed, ssd_pytorch_white_embed, faster_rcnn_pytorch_white_embed, \
-    classification_pytorch_white_embed, googlenet_pytorch_white_embed, classification_pytorch_black_embed, \
+    classification_pytorch_white_embed, googlenet_vgg16_pytorch_white_embed, classification_pytorch_black_embed, \
     classfication_tensorflow_white_embed, classfication_tensorflow_black_embed
 
 generator = Blueprint('generator', __name__)
@@ -98,10 +98,10 @@ def watermark_embed():
             ssd_pytorch_black_embed.modify_model_project(secret_label, extract_to_path, public_key)
         if model_value == 'ssd' and embed_type == 'whitebox':
             ssd_pytorch_white_embed.modify_model_project(secret_label, extract_to_path, public_key)
-        if (model_value in ['alexnet', 'vggnet', 'resnet']) and embed_type == 'whitebox':
+        if model_value in ['alexnet', 'resnet'] and embed_type == 'whitebox':
             classification_pytorch_white_embed.modify_model_project(secret_label, extract_to_path, public_key)
-        if model_value == 'googlenet' and embed_type == 'whitebox':
-            googlenet_pytorch_white_embed.modify_model_project(secret_label, extract_to_path, public_key)
+        if model_value in ['googlenet', 'vggnet'] and embed_type == 'whitebox':
+            googlenet_vgg16_pytorch_white_embed.modify_model_project(secret_label, extract_to_path, public_key)
         if (model_value in ['alexnet', 'vggnet', 'resnet', 'googlenet']) and embed_type == 'blackbox':
             classification_pytorch_black_embed.modify_model_project(secret_label, extract_to_path, public_key)
     # 压缩修改后的模型文件代码

watermark_generate/deals/classfication_tensorflow_white_embed.py

@@ -356,7 +356,7 @@ f"""def train_model(args, train_generator, val_generator):
         if isinstance(layer, tf.keras.layers.Conv2D):
             embed_layers.append(layer)
     # 使用指定的卷积层初始化
-    encoder = ModelEncoder(embed_layers[0:3], secret, "keys/key.npy", model)
+    encoder = ModelEncoder(embed_layers[3:7], secret, "keys/key.npy", model)
 
     # 调整学习率
     learning_rate = args.lr if args.lr else 1e-2

watermark_generate/deals/googlenet_pytorch_white_embed.py

@@ -1,5 +1,5 @@
 """
-GoogleNet 白盒水印嵌入工程文件(pytorch)处理
+GoogleNet, VGG16 白盒水印嵌入工程文件(pytorch)处理
 """
 import os