-
Notifications
You must be signed in to change notification settings - Fork 4
/
compute_jacobian.py
200 lines (179 loc) · 8.47 KB
/
compute_jacobian.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
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
# python3.7
"""Functions to compute Jacobian based on pre-trained GAN generator.
Support StyleGAN2 or StyleGAN3
"""
import os
import argparse
import warnings
from tqdm import tqdm
import numpy as np
import torch
import torch.nn.functional as F
from torch.autograd.functional import jacobian
from models import build_model
from utils.image_utils import save_image
from utils.image_utils import postprocess_image
from utils.custom_utils import to_numpy
warnings.filterwarnings(action='ignore', category=UserWarning)
def parse_args():
"""Parses arguments."""
parser = argparse.ArgumentParser()
group = parser.add_argument_group('General options.')
group.add_argument('weight_path', type=str,
help='Weight path to the pre-trained model.')
group.add_argument('--save_dir', type=str, default=None,
help='Directory to save the results. If not specified, '
'the results will be saved to '
'`work_dirs/{TASK_SPECIFIC}/` by default.')
group.add_argument('--job', type=str, default='jacobians',
help='Name for the job (default: jacobians)')
group.add_argument('--seed', type=int, default=4,
help='Seed for sampling. (default: 4)')
group.add_argument('--nums', type=int, default=5,
help='Number of samples to synthesized. (default: 5)')
group.add_argument('--img_size', type=int, default=1024,
help='Size of the synthesized images. (default: 1024)')
group.add_argument('--w_dim', type=int, default=512,
help='Dimension of the latent w. (default: 512)')
group.add_argument('--save_jpg', action='store_false',
help='Whether to save the images used to compute '
'jacobians. (default: True)')
group.add_argument('-d', '--data_name', type=str, default='ffhq',
help='Name of the datasets. (default: ffhq)')
group.add_argument('--latent_path', type=str, default='',
help='Path to the given latent codes. (default: None)')
group = parser.add_argument_group('StyleGAN2')
group.add_argument('--stylegan2', action='store_true',
help='Whether or not using StyleGAN2. (default: False)')
group.add_argument('--scale_stylegan2', type=float, default=1.0,
help='Scale for the number of channel fro stylegan2.')
group.add_argument('--randomize_noise', type=str, default='const',
help='Noise type when computing. (const or random)')
group = parser.add_argument_group('StyleGAN3')
group.add_argument('--stylegan3', action='store_true',
help='Whether or not using StyleGAN3. (default: False)')
group.add_argument('--cfg', type=str, default='T',
help='Config of the stylegan3 (T/R).')
group.add_argument('--scale_stylegan3r', type=float, default=2.0,
help='Scale for the number of channel for stylegan3 R.')
group.add_argument('--scale_stylegan3t', type=float, default=1.0,
help='Scale for the number of channel for stylegan3 T.')
group.add_argument('--tx', type=float, default=0,
help='Translate X-coordinate. (default: 0.0)')
group.add_argument('--ty', type=float, default=0,
help='Translate Y-coordinate. (default: 0.0)')
group.add_argument('--rotate', type=float, default=0,
help='Rotation angle in degrees. (default: 0)')
group = parser.add_argument_group('Jacobians')
group.add_argument('--b', type=float, default=1e-3,
help='Constant when computing jacobians fast.')
group.add_argument('--batch_size', type=int, default=4,
help='Batch size. (default: 4)')
return parser.parse_args()
def main():
"""Main function."""
args = parse_args()
# Parse model configuration.
assert (args.stylegan2 and not args.stylegan3) or \
(not args.stylegan2 and args.stylegan3)
job_disc = ''
if args.stylegan2:
config = dict(model_type='StyleGAN2Generator',
resolution=args.img_size,
w_dim=args.w_dim,
fmaps_base=int(args.scale_stylegan2 * (32 << 10)),
fmaps_max=512,)
job_disc += 'stylegan2'
else:
if args.stylegan3 and args.cfg == 'R':
config = dict(model_type='StyleGAN3Generator',
resolution=args.img_size,
w_dim=args.w_dim,
fmaps_base=int(args.scale_stylegan3r * (32 << 10)),
fmaps_max=1024,
use_radial_filter=True,)
job_disc += 'stylegan3r'
elif args.stylegan3 and args.cfg == 'T':
config = dict(model_type='StyleGAN3Generator',
resolution=args.img_size,
w_dim=args.w_dim,
fmaps_base=int(args.scale_stylegan3t * (32 << 10)),
fmaps_max=512,
use_radial_filter=False,
kernel_size=3,)
job_disc += 'stylegan3t'
else:
raise TypeError(f'StyleGAN3 config type error, need `R/T`,'
f' but got {args.cfg}')
job_name = f'seed_{args.seed}_num_{args.nums}_{job_disc}'
temp_dir = f'work_dirs/{args.job}/{args.data_name}/{job_name}'
save_dir = args.save_dir or temp_dir
os.makedirs(save_dir, exist_ok=True)
if args.save_jpg:
os.makedirs(f'{save_dir}/images', exist_ok=True)
print('Building generator...')
generator = build_model(**config)
checkpoint_path = args.weight_path
print(f'Loading checkpoint from `{checkpoint_path}` ...')
checkpoint = torch.load(checkpoint_path, map_location='cpu')['models']
if 'generator_smooth' in checkpoint:
generator.load_state_dict(checkpoint['generator_smooth'])
else:
generator.load_state_dict(checkpoint['generator'])
generator = generator.eval().cuda()
print('Finish loading checkpoint.')
# Set random seed.
np.random.seed(args.seed)
torch.manual_seed(args.seed)
if os.path.exists(args.latent_path):
latent_zs = np.load(args.latent_path)
latent_zs = latent_zs[:args.nums]
else:
latent_zs = np.random.randn(args.nums, generator.z_dim)
latent_zs = torch.from_numpy(latent_zs.astype(np.float32))
latent_zs = latent_zs.cuda()
with torch.no_grad():
latent_ws = generator.mapping(latent_zs)['w']
print(f'Shape of the latent w: {latent_ws.shape}')
def syn2jaco(w):
"""Wrap the synthesized function to compute the Jacobian easily.
Basically, this function defines a generator that takes the input
from the W space and then synthesizes an image. If the image is
larger than 256, it will be resized to 256 to save the time and
storage.
Args:
w: latent code from the W space
Returns:
An image with the size of [1, 256, 256]
"""
wp = w.unsqueeze(1).repeat((1, generator.num_layers, 1))
image = generator.synthesis(wp)['image']
if image.shape[-1] > 256:
scale = 256 / image.shape[-1]
image = F.interpolate(image, scale_factor=scale)
image = torch.sum(image, dim=1)
return image
jacobians = []
for idx in tqdm(range(latent_zs.shape[0])):
latent_w = latent_ws[idx:idx+1]
jac_i = jacobian(func=syn2jaco,
inputs=latent_w,
create_graph=False,
strict=False)
jacobians.append(jac_i)
if args.save_jpg:
wp = latent_w.unsqueeze(1).repeat((1, generator.num_layers, 1))
syn_outputs = generator.synthesis(wp)['image']
syn_outputs = to_numpy(syn_outputs)
images = postprocess_image(syn_outputs)
save_path = f'{save_dir}/images/{idx:06d}.jpg'
save_image(save_path, images[0])
jacobians = torch.cat(jacobians, dim=0)
jacobians = to_numpy(jacobians)
print(f'shape of the jacobian: {jacobians.shape}')
latent_ws = to_numpy(latent_ws)
np.save(f'{save_dir}/latent_codes.npy', latent_ws)
np.save(f'{save_dir}/jacobians_w.npy', jacobians)
print(f'Finish computing {args.nums} jacobians.')
if __name__ == '__main__':
main()