-
Notifications
You must be signed in to change notification settings - Fork 0
/
main.py
179 lines (156 loc) · 8.6 KB
/
main.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
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
# !/usr/bin/python3
# -*- coding:utf-8 -*-
# Author: Li JIANG
"""
PaddlePaddle Implementation of StarGAN-v2
"""
import os
import argparse
import sys
from munch import Munch
from core.data_reader import get_train_loader
from core.data_reader import get_test_loader
from core.solver import Solver
import paddle
import paddle.fluid as fluid
def str2bool(v):
return v.lower() in ('true')
def subdirs(dname):
return [d for d in os.listdir(dname)
if os.path.isdir(os.path.join(dname, d))]
def main(args):
print(args)
solver = Solver(args)
if args.mode == 'train':
assert len(subdirs(args.train_img_dir)) == args.num_domains
assert len(subdirs(args.val_img_dir)) == args.num_domains
loaders = Munch(src=get_train_loader(root=args.train_img_dir,
which='source',
img_size=args.img_size,
batch_size=args.batch_size,
prob=args.randcrop_prob,
num_workers=args.num_workers),
ref=get_train_loader(root=args.train_img_dir,
which='reference',
img_size=args.img_size,
batch_size=args.batch_size,
prob=args.randcrop_prob,
num_workers=args.num_workers),
val=get_test_loader(root=args.val_img_dir,
img_size=args.img_size,
batch_size=args.val_batch_size,
shuffle=True,
num_workers=args.num_workers))
solver.train(loaders)
elif args.mode == 'sample':
assert len(subdirs(args.src_dir)) == args.num_domains
assert len(subdirs(args.ref_dir)) == args.num_domains
loaders = Munch(src=get_test_loader(root=args.src_dir,
img_size=args.img_size,
batch_size=args.val_batch_size,
shuffle=False,
num_workers=args.num_workers),
ref=get_test_loader(root=args.ref_dir,
img_size=args.img_size,
batch_size=args.val_batch_size,
shuffle=False,
num_workers=args.num_workers))
solver.sample(loaders)
elif args.mode == 'eval':
solver.evaluate()
else:
raise NotImplementedError
if __name__ == '__main__':
parser = argparse.ArgumentParser()
# model arguments
parser.add_argument('--img_size', type=int, default=256,
help='Image resolution')
parser.add_argument('--num_domains', type=int, default=2,
help='Number of domains')
parser.add_argument('--latent_dim', type=int, default=8, # 16
help='Latent vector dimension')
parser.add_argument('--hidden_dim', type=int, default=64, # 512
help='Hidden dimension of mapping network')
parser.add_argument('--style_dim', type=int, default=8, # 64
help='Style code dimension')
# weight for objective functions
parser.add_argument('--lambda_reg', type=float, default=1,
help='Weight for R1 regularization')
parser.add_argument('--lambda_cyc', type=float, default=1,
help='Weight for cyclic consistency loss')
parser.add_argument('--lambda_sty', type=float, default=1,
help='Weight for style reconstruction loss')
parser.add_argument('--lambda_ds', type=float, default=1,
help='Weight for diversity sensitive loss')
parser.add_argument('--ds_iter', type=int, default=100000,
help='Number of iterations to optimize diversity sensitive loss')
parser.add_argument('--w_hpf', type=float, default=0,
help='weight for high-pass filtering') # default=1
# training arguments
parser.add_argument('--randcrop_prob', type=float, default=0.5,
help='Probabilty of using random-resized cropping')
parser.add_argument('--total_iters', type=int, default=2, # 100000
help='Number of total iterations')
parser.add_argument('--resume_iter', type=int, default=0,
help='Iterations to resume training/testing')
parser.add_argument('--batch_size', type=int, default=4, # 8
help='Batch size for training')
parser.add_argument('--val_batch_size', type=int, default=4, # 32
help='Batch size for validation')
parser.add_argument('--lr', type=float, default=1e-4,
help='Learning rate for D, E and G')
parser.add_argument('--f_lr', type=float, default=1e-6,
help='Learning rate for F')
parser.add_argument('--beta1', type=float, default=0.0,
help='Decay rate for 1st moment of Adam')
parser.add_argument('--beta2', type=float, default=0.99,
help='Decay rate for 2nd moment of Adam')
parser.add_argument('--weight_decay', type=float, default=1e-4,
help='Weight decay for optimizer')
parser.add_argument('--num_outs_per_domain', type=int, default=2, # 10
help='Number of generated images per domain during sampling')
# misc
parser.add_argument('--mode', type=str, default='train',
choices=['train', 'sample', 'eval', 'align'],
help='This argument is used in solver')
parser.add_argument('--num_workers', type=int, default=4,
help='Number of workers used in DataLoader')
parser.add_argument('--seed', type=int, default=777,
help='Seed for random number generator')
# directory for training
parser.add_argument('--train_img_dir', type=str, default='data_dev/celeba_hq/train',
help='Directory containing training images')
parser.add_argument('--val_img_dir', type=str, default='data_dev/celeba_hq/val',
help='Directory containing validation images')
parser.add_argument('--sample_dir', type=str, default='expr/samples',
help='Directory for saving generated images')
parser.add_argument('--checkpoint_dir', type=str, default='expr/checkpoints',
help='Directory for saving network checkpoints')
# directory for calculating metrics
parser.add_argument('--eval_dir', type=str, default='expr/eval',
help='Directory for saving metrics, i.e., FID and LPIPS')
# directory for testing
parser.add_argument('--result_dir', type=str, default='expr/results',
help='Directory for saving generated images and videos')
parser.add_argument('--src_dir', type=str, default='assets/representative/celeba_hq/src',
help='Directory containing input source images')
parser.add_argument('--ref_dir', type=str, default='assets/representative/celeba_hq/ref',
help='Directory containing input reference images')
parser.add_argument('--inp_dir', type=str, default='assets/representative/custom/female',
help='input directory when aligning faces')
parser.add_argument('--out_dir', type=str, default='assets/representative/celeba_hq/src/female',
help='output directory when aligning faces')
# face alignment
parser.add_argument('--wing_path', type=str, default='expr/checkpoints/wing.ckpt')
parser.add_argument('--lm_path', type=str, default='expr/checkpoints/celeba_lm_mean.npz')
# step size
parser.add_argument('--print_every', type=int, default=1)
parser.add_argument('--sample_every', type=int, default=1)
parser.add_argument('--save_every', type=int, default=1)
parser.add_argument('--eval_every', type=int, default=50000)
# choose gpu
parser.add_argument('--whichgpu', type=int, default=1)
args = parser.parse_args()
place = paddle.fluid.CUDAPlace(args.whichgpu) if paddle.fluid.is_compiled_with_cuda() else paddle.fluid.CPUPlace()
with fluid.dygraph.guard(place):
main(args)