Alexnet模型添加选择嵌入层方法

model/Alexnet.py

@@ -58,6 +58,16 @@ class Alexnet(nn.Module):
         x = self.classifier(x)
         return x
 
+    def get_encode_layers(self):
+        """
+        获取用于白盒模型水印加密层，每个模型根据复杂度选择合适的卷积层
+        """
+        conv_list = []
+        for module in self.modules():
+            if isinstance(module, nn.Conv2d):
+                conv_list.append(module)
+        return conv_list[0:2]
+
 if __name__ == '__main__':
     import argparse