forked from HongyangGao/PixelTCN
-
Notifications
You must be signed in to change notification settings - Fork 0
/
main.py.sample
50 lines (43 loc) · 2.02 KB
/
main.py.sample
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
import os
import time
import tensorflow as tf
from network import DilatedPixelCNN
def configure():
# training
flags = tf.app.flags
flags.DEFINE_integer('max_epoch', 3, '# of step in an epoch')
flags.DEFINE_integer('test_step', 100, '# of step to test a model')
flags.DEFINE_integer('save_step', 100, '# of step to save a model')
flags.DEFINE_integer('summary_step', 1, '# of step to save the summary')
flags.DEFINE_float('learning_rate', 1e-3, 'learning rate')
flags.DEFINE_float('keep_prob', 0.9, 'dropout probability')
flags.DEFINE_boolean('use_gpu', False, 'use GPU or not')
# data
flags.DEFINE_string('data_dir', './dataset', 'Name of data directory')
flags.DEFINE_string('train_list', './dataset/train.txt', 'Training data')
flags.DEFINE_string('valid_list', './dataset/val.txt', 'Training data')
flags.DEFINE_string('sample_dir', 'samples', 'Sample directory')
flags.DEFINE_integer('batch', 10, 'batch size')
flags.DEFINE_integer('channel', 3, 'channel size')
flags.DEFINE_integer('height', 256, 'height size')
flags.DEFINE_integer('width', 256, 'width size')
# Debug
flags.DEFINE_boolean('is_train', True, 'Training or testing')
flags.DEFINE_string('log_level', 'INFO', 'Log level')
flags.DEFINE_string('logdir', './logdir', 'Log dir')
flags.DEFINE_string('modeldir', './modeldir', 'Model dir')
flags.DEFINE_string('model_name', 'model', 'Model file name')
flags.DEFINE_integer('reload_step', 0, 'Reload step')
flags.DEFINE_integer('random_seed', int(time.time()), 'random seed')
# network
flags.DEFINE_integer('network_depth', 5, 'network depth for U-Net')
flags.DEFINE_integer('class_num', 21, 'output class number')
flags.DEFINE_integer('start_channel_num', 64, 'start number of outputs')
flags.FLAGS.__dict__['__parsed'] = False
return flags.FLAGS
def main(_):
model = DilatedPixelCNN(tf.Session(), configure())
model.train()
if __name__ == '__main__':
os.environ['CUDA_VISIBLE_DEVICES'] = '0'
tf.app.run()