|
|
@@ -86,10 +86,20 @@ class GoogLeNet(nn.Module):
|
|
|
x = self.a5(x)
|
|
|
x = self.b5(x)
|
|
|
x = self.avgpool(x)
|
|
|
- x = x.view(x.size(0), -1)
|
|
|
+ x = x.reshape(x.size(0), -1)
|
|
|
x = self.linear(x)
|
|
|
return x
|
|
|
|
|
|
+ def get_encode_layers(self):
|
|
|
+ """
|
|
|
+ 获取用于白盒模型水印加密层，每个模型根据复杂度选择合适的卷积层
|
|
|
+ """
|
|
|
+ conv_list = []
|
|
|
+ for module in self.modules():
|
|
|
+ if isinstance(module, nn.Conv2d) and module.out_channels > 100:
|
|
|
+ conv_list.append(module)
|
|
|
+ return conv_list[0:2]
|
|
|
+
|
|
|
if __name__ == '__main__':
|
|
|
import argparse
|
|
|
|
