-
Notifications
You must be signed in to change notification settings - Fork 0
/
main.py
85 lines (75 loc) · 3.4 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
#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Created on Tue Dec 12 14:18:55 2017
@author: jm
"""
import argparse
import os
import tensorflow as tf
from StarGAN_Model import stargan
# argument parser
parser = argparse.ArgumentParser(description='')
parser.add_argument('--phase', type=str, default='train', help='train or test')
parser.add_argument('--gpu_number', type=str, default='0')
parser.add_argument('--data_dir', type=str, default=os.path.join('.','data_mcep'))
parser.add_argument('--log_dir', type=str, default='log') # in assets/ directory
parser.add_argument('--ckpt_dir', type=str, default='checkpoint') # in assets/ directory
parser.add_argument('--sample_dir', type=str, default='sample') # in assets/ directory
parser.add_argument('--test_dir', type=str, default='test') # in assets/ directory
parser.add_argument('--epoch', type=int, default=100)
parser.add_argument('--batch_size', type=int, default=32)
parser.add_argument('--image_size', type=int, default=64)
parser.add_argument('--image_channel', type=int, default=1)
parser.add_argument('--nf', type=int, default=64) # number of filters
parser.add_argument('--n_label', type=int, default=7)
parser.add_argument('--lambda_gp', type=int, default=1)
parser.add_argument('--lambda_cls', type=int, default=1)
parser.add_argument('--lambda_rec', type=int, default=10)
parser.add_argument('--lambda_feat', type=int, default=2.5)
parser.add_argument('--lr', type=float, default=0.00001) # learning_rate
parser.add_argument('--lr_g', type=float, default=0.000015) # learning_rate
parser.add_argument('--lr_d', type=float, default=0.000015) # learning_rate
parser.add_argument('--beta1', type=float, default=0.5)
parser.add_argument('--continue_train', type=bool, default=False)
parser.add_argument('--snapshot', type=int, default=1286) # number of iterations to save files
parser.add_argument('--adv_type', type=str, default='GAN', help='GAN or WGAN')
parser.add_argument('--binary_attrs', type=str, default='0110000')
#self.attr_keys = ['Male', 'Female', 'KizunaAI', 'Nekomasu', 'Mirai', 'Shiro', 'Kaguya']
args = parser.parse_args()
def main(_):
# setting
os.environ["CUDA_DEVICE_ORDER"] = "PCI_BUS_ID" # see issue #152
os.environ["CUDA_VISIBLE_DEVICES"] = args.gpu_number
tf.reset_default_graph()
assets_dir = os.path.join('.', 'assets', '{}_label{}_img{}'.format(args.adv_type, args.n_label, args.image_size))
args.log_dir = os.path.join(assets_dir, args.log_dir)
args.ckpt_dir = os.path.join(assets_dir, args.ckpt_dir)
args.sample_dir = os.path.join(assets_dir, args.sample_dir)
args.test_dir = os.path.join(assets_dir, args.test_dir)
# make directory if not exist
try:
os.makedirs(args.log_dir)
except:
pass
try:
os.makedirs(args.ckpt_dir)
except:
pass
try:
os.makedirs(args.sample_dir)
except:
pass
try:
os.makedirs(args.test_dir)
except:
pass
# run session
tfconfig = tf.ConfigProto()
tfconfig.gpu_options.allow_growth = True
with tf.Session(config=tfconfig) as sess:
model = stargan(sess, args)
model.train() if args.phase == 'train' else model.test()
# run main function
if __name__ == '__main__':
tf.app.run()