Skip to content

Latest commit

 

History

History
95 lines (66 loc) · 4.95 KB

README.md

File metadata and controls

95 lines (66 loc) · 4.95 KB
Raw

Watermark-DnCNN

Intro

This repo refers to Watermarking Dnns paper you can find here.

Requirements

In order to execute this code, following programs are required:

  • Python 3.7
  • tensorflow 1.15 (tested, but can be either higher or slithy slower)
  • Pillow, OpenCV, Matplotlib, Numpy (for windows use max numpy version 1.19.3 becayse 1.19.4 will fail sanity check)

It is advised to install tensorflow in an environment (venv or conda) for better management.

Folders

Below scripts saves training checkpoints in various folders. Those directories are used:

  • DnCNN_weight: contains dncnn's weights after been trained for 45 epochs. This is used as 'basic model' like the customer's model that needs to be watermarked.
  • overwriting: contains checkpoints from retraining base dncnn model (in DnCNN_weight) for watermarking.
  • combine_weight: contains checkpoints from retraining other models. Used as Deep Prior model (Auxiliary Visualizer).

Fast start

In order to test the repo, execute:

  1. GeneratorTriggerVerificationImg.py that will create keys folder with Trigger and Verification images
  2. Preprocess_dataset_for_dncnn.py that will create all needed datas
  3. Watermark_train.py in order to train dncnn for 8 epochs (ckpt saved in overwriting)
  4. AuxVisualizer_train.py in order to train Prior model for 8 epochs (ckpt saved in combine_weight)
  5. ExecuteFineTuning.py in order to finetune watermarked model (ckpt saved in fineTuning_weight)
  6. ExecutePruning.py in order to prune watermarked model (ckpt saved in pruning_weights)
  7. PaperAnalysis reproduces the uniqueness, fidelity and robustness tests and shows interesting plots

** note: the scripts above, in 1, 2, 3, 4 can be run directly via the fastrun_train.py

Scripts description

Following python files are only used as import modules for other scripts:

  • DnCNN_model.py: contains all functions needed to create the dncnn model (also with loss and optimizer). If it is runned, it will compile and run once the dncnn printing all layers and testing if allright

  • AuxVisualizerModel.py: contains all functions needed to create Deep Prior model (also with loss and optimizer)

  • utility.py: contains utility functions for file store and other utilities

  • WatermarkedTrainedModel.py nad WatermarkedVisualizerModel: easy to use wrappers for needed modules for Paper tests All the scripts can be logically splitted into those categories:

  • prunings.py: contains utility functions for weight and unit pruning on tensorflow Variable

  • Preprocessors: DnCNN_test.py for basic environment test, GeneratorTriggerVerificationImg.py for trigger/verification keys, Preprocess_dataset_for_dncnn.py for numpy datasets creation

  • Fast run scripts: fastrun_train.py and fastrun_test.py

  • Watermarking training: DnCNN_model.py (model lib), Watermark_train.py for embedding watermark

  • Auxiliary visualizer: AuxVisualizerModel.py (model lib), AuxVisualizer_train.py for training 8 epochs

  • Model storing library: WatermarkedTrainedModel.py (fast loading class to manage a model from checkpoint)

  • Principles related scripts: ExecuteFineTuning.py for Finetuning on 100 epochs (on Img12 or KTH dataset), ExecutePruning.py for pruning with k% weights, prunings.py (pruning utility lib)

  • Paper results checking: PaperAnalysis.py

Approaces tested thru PaperAnalysis

Script PaperAnalysis.py tests given dncnn model on:

  • Uniqueness
  • Robustness (finetuning)
  • Robustness (pruning)
  • Watermarking vs Original model output
  • Fidelity

Just set boolean True on interesting variables in PaperAnalysis main's method and run it. When show_images parameter (from various methods) is set to True, remember to press space on image to go further (opencv imshow method manages it that way). Results are saved in results subfolders with images and json datas. Below are some visual simple results tested.

Results

All the results can be viewed in attached paper (TBD).

Watermarking the dnn

Finetuning all layers with trigger/verification key results in watermarking embedding:

Robustness against Finetuning

When finetuning last convolutional layer, the watermark is keep even for first 100 epochs.

Robustness against Pruning

When Pruning all convolutional layers, the watermark is kept until k=60%.



Original code: painfulloop (Main paper author)

Paper results implementation contributors: Fabian Greavu and Scaramuzzino Giovanna (scripts and results added within Image Processing & Security course project at the Department of Information Engineering of the University of Florence under Prof. Alessandro Piva's supervision)