update eval

caradryanl · May 15, 2024 · 50e8f15 · 50e8f15
1 parent cba88f6
commit 50e8f15
Show file tree

Hide file tree

Showing 3 changed files with 68 additions and 32 deletions.
diff --git a/diffusers/scripts/train_pfami.py b/diffusers/scripts/train_pfami.py
@@ -13,7 +13,7 @@
 import time, json
 
 from stable_copyright import PFAMIStableDiffusionPipeline, SecMIDDIMScheduler, PFAMILatentDiffusionPipeline
-from stable_copyright import load_dataset, benchmark
+from stable_copyright import load_dataset, benchmark, test
 
 def image_perturbation(image, strength, image_size=512):
     perturbation = transforms.Compose([
@@ -148,22 +148,32 @@ def main(args):
         torch.save(nonmember_scores_all_steps, args.output + f'pfami_{args.model_type}_nonmember_scores_all_steps.pth')
 
         member_corr_scores, nonmember_corr_scores = compute_corr_score(member_scores_all_steps, nonmember_scores_all_steps)
-
-        benchmark(member_scores_sum_step, nonmember_scores_sum_step, f'pfami_{args.model_type}_sum_score', args.output)
-        benchmark(member_corr_scores, nonmember_corr_scores, f'pfami_{args.model_type}_corr_score', args.output)
 
-        with open(args.output + f'pfami_{args.model_type}_image_log.json', 'w') as file:
-            json.dump(dict(member=member_path_log, nonmember=nonmember_path_log), file, indent=4)
+        if not args.eval:
+            benchmark(member_scores_sum_step, nonmember_scores_sum_step, f'pfami_{args.model_type}_sum_score', args.output)
+            benchmark(member_corr_scores, nonmember_corr_scores, f'pfami_{args.model_type}_corr_score', args.output)
+
+            with open(args.output + f'pfami_{args.model_type}_image_log.json', 'w') as file:
+                json.dump(dict(member=member_path_log, nonmember=nonmember_path_log), file, indent=4)
+
+            end_time = time.time()
+            elapsed_time = end_time - start_time
+            running_time = dict(running_time=elapsed_time)
+            with open(args.output + f'pfami_{args.model_type}_running_time.json', 'w') as file:
+                json.dump(running_time, file, indent=4)
+        else:
+            threshold_path = args.threshold_root + f'{args.model_type}/pfami/'
+
+            test(member_scores_sum_step, member_scores_sum_step, f'pfami_{args.model_type}_sum_score', args.output, threshold_path)
+            test(member_corr_scores, nonmember_corr_scores, f'pfami_{args.model_type}_corr_score', args.output, threshold_path)
+
+            with open(args.output + f'pfami_{args.model_type}_image_log_test.json', 'w') as file:
+                json.dump(dict(member=member_path_log, nonmember=nonmember_path_log), file, indent=4)
 
     else:
         raise NotImplementedError('DDP not implemented')
 
-    end_time = time.time()
-    elapsed_time = end_time - start_time
-    running_time = dict(running_time=elapsed_time)
 
-    with open(args.output + f'pfami_{args.model_type}_running_time.json', 'w') as file:
-        json.dump(running_time, file, indent=4)
 
 
 
@@ -194,6 +204,8 @@ def fix_seed(seed):
     parser.add_argument('--end-strength', type=float, default=0.7)
     parser.add_argument('--model-type', type=str, choices=['sd', 'sdxl', 'ldm'], default='sd')
     parser.add_argument('--demo', type=bool, default=False)
+    parser.add_argument('--eval', type=bool, default=False)
+    parser.add_argument('--threshold-root', type=str, default='experiments/')
     args = parser.parse_args()
 
     fix_seed(args.seed)

diff --git a/diffusers/scripts/train_pia.py b/diffusers/scripts/train_pia.py
@@ -10,7 +10,7 @@
 import json,time
 
 from stable_copyright import PIAStableDiffusionPipeline, SecMIDDIMScheduler, PIALatentDiffusionPipeline
-from stable_copyright import load_dataset, benchmark
+from stable_copyright import load_dataset, benchmark, test
 
 
 def load_pipeline(ckpt_path, device='cuda:0', model_type='sd'):
@@ -110,21 +110,32 @@ def main(args):
 
         member_corr_scores, nonmember_corr_scores = compute_corr_score(member_scores_all_steps, nonmember_scores_all_steps)
 
-        benchmark(member_scores_sum_step, nonmember_scores_sum_step, f'{pia_or_pian}_{args.model_type}_sum_score', args.output)
-        benchmark(member_corr_scores, nonmember_corr_scores, f'{pia_or_pian}_{args.model_type}_corr_score', args.output)
+
+        if not args.eval:
+            benchmark(member_scores_sum_step, nonmember_scores_sum_step, f'{pia_or_pian}_{args.model_type}_sum_score', args.output)
+            benchmark(member_corr_scores, nonmember_corr_scores, f'{pia_or_pian}_{args.model_type}_corr_score', args.output)
+
+            with open(args.output + f'{pia_or_pian}_{args.model_type}_image_log.json', 'w') as file:
+                json.dump(dict(member=member_path_log, nonmember=nonmember_path_log), file, indent=4)
+
+            end_time = time.time()
+            elapsed_time = end_time - start_time
+            running_time = dict(running_time=elapsed_time)
+            with open(args.output + f'{pia_or_pian}_{args.model_type}_running_time.json', 'w') as file:
+                json.dump(running_time, file, indent=4)
+        else:
+            threshold_path = args.threshold_root + f'{args.model_type}/{pia_or_pian}/'
+
+            test(member_scores_sum_step, member_scores_sum_step, f'{pia_or_pian}_{args.model_type}_sum_score', args.output, threshold_path)
+            test(member_corr_scores, nonmember_corr_scores, f'{pia_or_pian}_{args.model_type}_corr_score', args.output, threshold_path)
 
-        with open(args.output + f'{pia_or_pian}_{args.model_type}_image_log.json', 'w') as file:
-            json.dump(dict(member=member_path_log, nonmember=nonmember_path_log), file, indent=4)
+            with open(args.output + f'{pia_or_pian}_{args.model_type}_image_log_test.json', 'w') as file:
+                json.dump(dict(member=member_path_log, nonmember=nonmember_path_log), file, indent=4)
 
     else:
         raise NotImplementedError('DDP not implemented')
 
-    end_time = time.time()
-    elapsed_time = end_time - start_time
-    running_time = dict(running_time=elapsed_time)
 
-    with open(args.output + f'{pia_or_pian}_{args.model_type}_running_time.json', 'w') as file:
-        json.dump(running_time, file, indent=4)
 
 
 
@@ -152,6 +163,8 @@ def fix_seed(seed):
     parser.add_argument('--normalized', type=bool, default=False)
     parser.add_argument('--model-type', type=str, choices=['sd', 'sdxl', 'ldm'], default='sd')
     parser.add_argument('--demo', type=bool, default=False)
+    parser.add_argument('--eval', type=bool, default=False)
+    parser.add_argument('--threshold-root', type=str, default='experiments/')
     args = parser.parse_args()
 
     fix_seed(args.seed)

diff --git a/diffusers/scripts/train_secmi.py b/diffusers/scripts/train_secmi.py
@@ -10,7 +10,7 @@
 import json,time
 
 from stable_copyright import SecMILatentDiffusionPipeline, SecMIStableDiffusionPipeline, SecMIDDIMScheduler
-from stable_copyright import load_dataset, benchmark
+from stable_copyright import load_dataset, benchmark, test
 
 
 def load_pipeline(ckpt_path, device='cuda:0', model_type='sd'):
@@ -101,21 +101,30 @@ def main(args):
 
         member_corr_scores, nonmember_corr_scores = compute_corr_score(member_scores_all_steps, nonmember_scores_all_steps)
 
-        benchmark(member_scores_50th_step, nonmember_scores_50th_step, f'secmi_{args.model_type}_50th_score', args.output)
-        benchmark(member_corr_scores, nonmember_corr_scores, f'secmi_{args.model_type}_corr_score', args.output)
-
-        with open(args.output + f'secmi_{args.model_type}_image_log.json', 'w') as file:
-            json.dump(dict(member=member_path_log, nonmember=nonmember_path_log), file, indent=4)
+        if not args.eval:
+            benchmark(member_scores_50th_step, nonmember_scores_50th_step, f'secmi_{args.model_type}_50th_score', args.output)
+            benchmark(member_corr_scores, nonmember_corr_scores, f'secmi_{args.model_type}_corr_score', args.output)
+
+            with open(args.output + f'secmi_{args.model_type}_image_log.json', 'w') as file:
+                json.dump(dict(member=member_path_log, nonmember=nonmember_path_log), file, indent=4)
+            end_time = time.time()
+            elapsed_time = end_time - start_time
+            running_time = dict(running_time=elapsed_time)
+            with open(args.output + f'secmi_{args.model_type}_running_time.json', 'w') as file:
+                json.dump(running_time, file, indent=4)
+        else:
+            threshold_path = args.threshold_root + f'{args.model_type}/secmi/'
+
+            test(member_scores_50th_step, nonmember_scores_50th_step, f'secmi_{args.model_type}_50th_score', args.output, threshold_path)
+            test(member_corr_scores, nonmember_corr_scores, f'secmi_{args.model_type}_corr_score', args.output, threshold_path)
+
+            with open(args.output + f'secmi_{args.model_type}_image_log_test.json', 'w') as file:
+                json.dump(dict(member=member_path_log, nonmember=nonmember_path_log), file, indent=4)
 
     else:
         raise NotImplementedError('DDP not implemented')
 
-    end_time = time.time()
-    elapsed_time = end_time - start_time
-    running_time = dict(running_time=elapsed_time)
 
-    with open(args.output + f'secmi_{args.model_type}_running_time.json', 'w') as file:
-        json.dump(running_time, file, indent=4)
 
 
 
@@ -142,6 +151,8 @@ def fix_seed(seed):
     parser.add_argument('--use-ddp', type=bool, default=False)
     parser.add_argument('--model-type', type=str, choices=['sd', 'sdxl', 'ldm'], default='sd')
     parser.add_argument('--demo', type=bool, default=False)
+    parser.add_argument('--eval', type=bool, default=False)
+    parser.add_argument('--threshold-root', type=str, default='experiments/')
     args = parser.parse_args()
 
     fix_seed(args.seed)