-
Notifications
You must be signed in to change notification settings - Fork 9
/
main_test_diff_passwds.py
149 lines (121 loc) · 5.27 KB
/
main_test_diff_passwds.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
# system libraries
import os, sys
import os.path as osp
import time
import numpy as np
from PIL import Image
import gc
from collections import OrderedDict
import torch
import torchvision.transforms as transforms
# libraries within this package
from cmd_args import parse_args
from utils.visualizer import Visualizer
from utils.util import generate_code
import datasets
import models
TEST_CODE_NUM = 10
def main():
# parse args
global args
args = parse_args(sys.argv[1])
args.during_training = False
args.gpu_ids = list(range(len(os.environ['CUDA_VISIBLE_DEVICES'].split(','))))
args.device = torch.device('cuda:0')
args.test_size = args.batch_size // 4 * len(args.gpu_ids)
# add timestamp to ckpt_dir
args.timestamp = time.strftime('%m%d%H%M%S', time.localtime())
args.ckpt_dir += '_' + args.timestamp
# -------------------- init ckpt_dir, logging --------------------
os.makedirs(args.ckpt_dir, mode=0o777, exist_ok=True)
# -------------------- init visu --------------------
visualizer = Visualizer(args)
visualizer.logger.log('sys.argv:\n' + ' '.join(sys.argv))
for arg in sorted(vars(args)):
visualizer.logger.log('{:20s} {}'.format(arg, getattr(args, arg)))
visualizer.logger.log('')
# -------------------- dataset & loader --------------------
test_dataset = datasets.__dict__[args.dataset](
train=False,
transform=transforms.Compose([
transforms.Resize(args.imageSize, Image.BICUBIC),
transforms.ToTensor(),
transforms.Normalize((0.5, 0.5, 0.5),
(0.5, 0.5, 0.5))
]),
args=args
)
visualizer.logger.log('test_dataset: ' + str(test_dataset))
test_loader = torch.utils.data.DataLoader(
test_dataset,
batch_size=args.test_size,
shuffle=False,
num_workers=args.workers,
pin_memory=True,
worker_init_fn=lambda x: np.random.seed((torch.initial_seed()) % (2 ** 32))
)
# -------------------- create model --------------------
model_dict = {}
G_input_nc = args.input_nc + args.passwd_length
model_dict['G'] = models.define_G(G_input_nc, args.output_nc,
args.ngf, args.which_model_netG, args.n_downsample_G,
args.normG, args.dropout,
args.init_type, args.init_gain,
args.passwd_length,
use_leaky=args.use_leakyG,
use_resize_conv=args.use_resize_conv,
padding_type=args.padding_type)
model_dict['G_nets'] = [model_dict['G']]
print('model_dict')
for k, v in model_dict.items():
print(k + ':')
if isinstance(v, list):
print('list, len:', len(v))
print('')
else:
print(v)
# -------------------- resume --------------------
if args.resume:
if osp.isfile(args.resume):
print("=> loading checkpoint '{}'".format(args.resume))
checkpoint = torch.load(args.resume, map_location='cpu')
args.start_epoch = checkpoint['epoch'] + 1
name = 'G'
net = model_dict[name]
if isinstance(net, torch.nn.DataParallel):
net = net.module
net.load_state_dict(checkpoint['state_dict_' + name])
print("=> loaded checkpoint '{}' (epoch {})"
.format(args.resume, checkpoint['epoch']))
else:
print("=> no checkpoint found at '{}'".format(args.resume))
gc.collect()
torch.cuda.empty_cache()
test(test_loader, model_dict, visualizer, args)
def test(test_loader, model_dict, visualizer, args, iter=0):
model_dict['G'].train()
with torch.no_grad():
for dis_idx in range(TEST_CODE_NUM):
z, dis_target, \
rand_z, rand_dis_target, \
inv_z, inv_dis_target, \
rand_inv_z, rand_inv_dis_target, _, _ = generate_code(args.passwd_length,
args.batch_size,
args.device,
inv=True,
use_minus_one=args.use_minus_one,
gen_random_WR=False)
# all the test images use the same passwords
for i, (img, label, landmarks, img_path) in enumerate(test_loader):
fake = model_dict['G'](img, z.cpu())
recon = model_dict['G'](fake, inv_z)
rand_recon = model_dict['G'](fake, rand_inv_z)
current_visuals = OrderedDict()
current_visuals['real'] = img
current_visuals['fake'] = fake
current_visuals['recon'] = recon
current_visuals['rand_recon'] = rand_recon
visualizer.display_test_results_vertical_html(current_visuals, img_path, dis_idx, iter, use_real=True,
add_padding=False, refresh=-1)
if __name__ == '__main__':
main()