From 50e8f15446d67b48dd83845c948a7c9d6d660cad Mon Sep 17 00:00:00 2001 From: CaradryanLiang Date: Wed, 15 May 2024 16:52:15 -0700 Subject: [PATCH] update eval --- diffusers/scripts/train_pfami.py | 34 +++++++++++++++++++++----------- diffusers/scripts/train_pia.py | 33 +++++++++++++++++++++---------- diffusers/scripts/train_secmi.py | 33 ++++++++++++++++++++----------- 3 files changed, 68 insertions(+), 32 deletions(-) diff --git a/diffusers/scripts/train_pfami.py b/diffusers/scripts/train_pfami.py index 7e33e22..1973260 100644 --- 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 index bee05ea..036db1d 100644 --- 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 index 4801695..a4fc933 100644 --- 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)