-
Notifications
You must be signed in to change notification settings - Fork 6
/
evaluate_models.py
178 lines (137 loc) · 6.21 KB
/
evaluate_models.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
#!/usr/bin/env python
# -*- coding: utf-8 -*-
"""
Main file to evaluate the models
"""
import argparse
import os
import tensorflow as tf
from torch.utils.data import DataLoader
from deep_adversarial_network.adversarial_training import DeepGAN
from deep_adversarial_network.discriminator import (get_available_discriminators, make_discriminator)
from deep_adversarial_network.generator import (get_available_generators, make_generator)
from deep_adversarial_network.logging.logger import rootLogger
from deep_adversarial_network.utils import (get_available_datasets,
make_dataset, RNG)
from deep_adversarial_network.utils.pytorch_dataset_utils import DatasetIndexer
# Optimizers
OPTIMIZERS = {
'adam': tf.train.AdamOptimizer,
'adagrad': tf.train.AdagradOptimizer,
'sgd': tf.train.GradientDescentOptimizer,
'rms_prop': tf.train.RMSPropOptimizer
}
# Losses
LOSSES = {
'l1': tf.losses.absolute_difference,
'l2': tf.losses.mean_squared_error
}
# General Paths
LOG_PATH = os.path.join(os.getcwd(), 'logs/')
PLOT_PATH = os.path.join(os.getcwd(), 'plots/')
MODEL_PATH = os.path.join(os.getcwd(), 'models/')
IMAGE_PATH = os.path.join(os.getcwd(), 'data/')
# training settings
parser = argparse.ArgumentParser(formatter_class=argparse.ArgumentDefaultsHelpFormatter)
# general
parser.add_argument('-d', '--dataset', type=str, default='big',
help="dataset, {'" + \
"', '".join(get_available_datasets()) + \
"'}")
parser.add_argument('--data-dirpath', type=str, default='data/',
help='directory for storing downloaded data')
parser.add_argument('--n-workers', type=int, default=2,
help='how many threads to use for I/O')
parser.add_argument('--gpu', type=str, default='0',
help="ID of the GPU to train on (or '' to train on CPU)")
parser.add_argument('-rs', '--random-seed', type=int, default=1,
help="random seed for training")
# GAN-related
parser.add_argument('-dr', '--discriminator', type=str, default='resnet',
help="discriminator architecture name, {'" + \
"', '".join(get_available_discriminators()) + \
"'}")
parser.add_argument('-gr', '--generator', type=str, default='skip2',
help="generator architecture name, {'" + \
"', '".join(get_available_generators()) + \
"'}")
parser.add_argument('-d_lr', '--d_lr', type=float, default=1e-4,
help='discriminator learning rate')
parser.add_argument('-g_lr', '--g_lr', type=float, default=1e-4,
help='generator learning rate')
parser.add_argument('-b', '--batch_size', type=int, default=2,
help='input batch size for training')
parser.add_argument('-d_opt', '--d_optim', type=str, default='adam',
help="optimizer, {'" + \
"', '".join(OPTIMIZERS.keys()) + \
"'}")
parser.add_argument('-g_opt', '--g_optim', type=str, default='adam',
help="optimizer, {'" + \
"', '".join(OPTIMIZERS.keys()) + \
"'}")
parser.add_argument('-m', '--model_name', type=str,
default='hue_sat_total', help='name for model')
parser.add_argument('-e', '--epochs', type=int, default=1000,
help='number of epochs')
parser.add_argument('-rl', '--recon_loss', type=str, default='l2',
help="losses, {'" + \
"', '".join(LOSSES.keys()) + \
"'}")
# Plot related
parser.add_argument('-tf', '--tf_logs', type=str, default='tf_logs',
help="log folder for tensorflow logging")
parser.add_argument('-mp', '--plot_matplotlib', type=str, default='n',
help="whether to plot matplotlib plots")
# parse and validate parameters
args = parser.parse_args()
for k, v in args._get_kwargs():
if isinstance(v, str):
setattr(args, k, v.strip().lower())
os.environ['CUDA_VISIBLE_DEVICES'] = args.gpu
# print arguments
rootLogger.info("Running with the following parameters:")
rootLogger.info(vars(args))
def main(args=args):
"""
main function that parses the arguments and trains
:param args: arguments related
:return: None
"""
# load and shuffle data
dataset = make_dataset(args.dataset)
_, val_dataset = dataset.load(args.data_dirpath)
rng = RNG(args.random_seed)
val_ind = rng.permutation(len(val_dataset))
val_dataset = DatasetIndexer(val_dataset, val_ind)
mplib = True if args.plot_matplotlib == 'y' else False
batch_size = args.batch_size
recon_loss = LOSSES.get(args.recon_loss, None)
if not recon_loss:
raise ValueError("invalid loss: '{0}'".format(args.recon_loss))
val_loader = DataLoader(dataset=val_dataset,
batch_size=batch_size,
shuffle=False,
num_workers=args.n_workers)
# build discriminator model
discriminator = make_discriminator(name=args.discriminator)
# build generator model
generator = make_generator(name=args.generator)
# get optimizer
d_optim = OPTIMIZERS.get(args.d_optim, None)
if not d_optim:
raise ValueError("invalid optimizer: '{0}'".format(args.d_optim))
g_optim = OPTIMIZERS.get(args.g_optim, None)
if not g_optim:
raise ValueError("invalid optimizer: '{0}'".format(args.g_optim))
# get learning rate
d_lr = args.d_lr
g_lr = args.g_lr
# Create GAN according to params
model = DeepGAN(discriminator=discriminator, generator=generator, model_name=args.model_name, recon_loss=recon_loss,
dataset=args.dataset, batch_size=args.batch_size, d_optim=d_optim, g_optim=g_optim, d_lr=d_lr,
g_lr=g_lr,epochs=args.epochs, mplib=mplib, tf_log_path="")
# Test the model
model.validate_results(test_loader=val_loader, model_path=MODEL_PATH, image_path_gt=IMAGE_PATH + "gt/", image_path_pred=IMAGE_PATH +"pred/",
image_path_comp=IMAGE_PATH + "comp/")
if __name__ == '__main__':
main()