-
Notifications
You must be signed in to change notification settings - Fork 16
/
run_metrics.py
113 lines (93 loc) · 5.58 KB
/
run_metrics.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
# Copyright (c) 2019, NVIDIA CORPORATION. All rights reserved.
#
# This work is licensed under the Creative Commons Attribution-NonCommercial
# 4.0 International License. To view a copy of this license, visit
# http://creativecommons.org/licenses/by-nc/4.0/ or send a letter to
# Creative Commons, PO Box 1866, Mountain View, CA 94042, USA.
"""Main entry point for training StyleGAN and ProGAN networks."""
import dnnlib
from dnnlib import EasyDict
import dnnlib.tflib as tflib
import config
from metrics_evaluation import metric_base
# from training import misc
import misc
#----------------------------------------------------------------------------
def run_pickle(submit_config, metric_args, network_pkl, dataset_args, mirror_augment):
ctx = dnnlib.RunContext(submit_config)
tflib.init_tf()
print('Evaluating %s metric on network_pkl "%s"...' % (metric_args.name, network_pkl))
metric = dnnlib.util.call_func_by_name(**metric_args)
print()
metric.run(network_pkl, dataset_args=dataset_args, mirror_augment=mirror_augment, num_gpus=submit_config.num_gpus)
print()
ctx.close()
#----------------------------------------------------------------------------
def run_snapshot(submit_config, metric_args, run_id, snapshot):
ctx = dnnlib.RunContext(submit_config)
tflib.init_tf()
print('Evaluating %s metric on run_id %s, snapshot %s...' % (metric_args.name, run_id, snapshot))
run_dir = misc.locate_run_dir(run_id)
network_pkl = misc.locate_network_pkl(run_dir, snapshot)
metric = dnnlib.util.call_func_by_name(**metric_args)
print()
metric.run(network_pkl, run_dir=run_dir, num_gpus=submit_config.num_gpus)
print()
ctx.close()
#----------------------------------------------------------------------------
def run_all_snapshots(submit_config, metric_args, run_id):
ctx = dnnlib.RunContext(submit_config)
tflib.init_tf()
print('Evaluating %s metric on all snapshots of run_id %s...' % (metric_args.name, run_id))
run_dir = misc.locate_run_dir(run_id)
network_pkls = misc.list_network_pkls(run_dir)
metric = dnnlib.util.call_func_by_name(**metric_args)
print()
for idx, network_pkl in enumerate(network_pkls):
ctx.update('', idx, len(network_pkls))
metric.run(network_pkl, run_dir=run_dir, num_gpus=submit_config.num_gpus)
print()
ctx.close()
#----------------------------------------------------------------------------
def main():
submit_config = dnnlib.SubmitConfig()
# Which metrics to evaluate?
metrics = []
metrics += [metric_base.fid50k]
#metrics += [metric_base.ppl_zfull]
#metrics += [metric_base.ppl_wfull]
#metrics += [metric_base.ppl_zend]
#metrics += [metric_base.ppl_wend]
#metrics += [metric_base.ls]
#metrics += [metric_base.dummy]
# Which networks to evaluate them on?
tasks = []
# tasks += [EasyDict(run_func_name='run_metrics.run_pickle', network_pkl='../model_results/stylegan/network-snapshot-008908.pkl', dataset_args=EasyDict(tfrecord_dir='ffhq_128x128_tf_stylegan', shuffle_mb=0), mirror_augment=True)] # karras2019stylegan-ffhq-1024x1024.pkl
tasks += [EasyDict(run_func_name='run_metrics.run_pickle', network_pkl='/home/hxz/GraduationProject/newBeholder-GAN/results/beholder-6001/network-final.pkl', dataset_args=EasyDict(tfrecord_dir='tf_files', shuffle_mb=0), mirror_augment=True)] # karras2019stylegan-ffhq-1024x1024.pkl
# tasks += [EasyDict(run_func_name='run_metrics.run_pickle', network_pkl='../model_results/stylegan/karras2019stylegan-ffhq-1024x1024.pkl', dataset_args=EasyDict(tfrecord_dir='ffhq', shuffle_mb=0), mirror_augment=True)] # karras2019stylegan-ffhq-1024x1024.pkl
#tasks += [EasyDict(run_func_name='run_metrics.run_pickle', network_pkl='../model_results/beholdergan-original/network-final.pkl', dataset_args=EasyDict(tfrecord_dir='ffhq_128x128_tf', shuffle_mb=0), mirror_augment=True)] # beholdergan
# tasks += [EasyDict(run_func_name='run_metrics.run_pickle', network_pkl='../model_results/beholdergan-id/first/network-final.pkl', dataset_args=EasyDict(tfrecord_dir='ffhq_128x128_id_tf', shuffle_mb=0), mirror_augment=True)] # beholdergan-id
# tasks += [EasyDict(run_func_name='run_metrics.run_pickle', network_pkl='../model_results/beholdergan-id/second/network-final.pkl', dataset_args=EasyDict(tfrecord_dir='ffhq_128x128_id_tf', shuffle_mb=0), mirror_augment=True)] # beholdergan-id
#tasks += [EasyDict(run_func_name='run_metrics.run_snapshot', run_id=100, snapshot=25000)]
#tasks += [EasyDict(run_func_name='run_metrics.run_all_snapshots', run_id=100)]
# How many GPUs to use?
submit_config.num_gpus = 1
#submit_config.num_gpus = 2
#submit_config.num_gpus = 4
#submit_config.num_gpus = 8
# Execute.
submit_config.run_dir_root = dnnlib.submission.submit.get_template_from_path(config.result_dir)
submit_config.run_dir_ignore += config.run_dir_ignore
for task in tasks:
for metric in metrics:
submit_config.run_desc = '%s-%s' % (task.run_func_name, metric.name)
if task.run_func_name.endswith('run_snapshot'):
submit_config.run_desc += '-%s-%s' % (task.run_id, task.snapshot)
if task.run_func_name.endswith('run_all_snapshots'):
submit_config.run_desc += '-%s' % task.run_id
submit_config.run_desc += '-%dgpu' % submit_config.num_gpus
dnnlib.submit_run(submit_config, metric_args=metric, **task)
#----------------------------------------------------------------------------
if __name__ == "__main__":
main()
#----------------------------------------------------------------------------