File size: 15,079 Bytes
ba5dcdc
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
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
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
# Copyright 2020 Erik Härkönen. All rights reserved.
# This file is licensed to you under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License. You may obtain a copy
# of the License at http://www.apache.org/licenses/LICENSE-2.0

# Unless required by applicable law or agreed to in writing, software distributed under
# the License is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR REPRESENTATIONS
# OF ANY KIND, either express or implied. See the License for the specific language
# governing permissions and limitations under the License.

# Patch for broken CTRL+C handler
# https://github.com/ContinuumIO/anaconda-issues/issues/905
import os
os.environ['FOR_DISABLE_CONSOLE_CTRL_HANDLER'] = '1'

import numpy as np
import os
from pathlib import Path
import re
import sys
import datetime
import argparse
import torch
import json
from types import SimpleNamespace
import scipy
from scipy.cluster.vq import kmeans
from tqdm import trange
from netdissect.nethook import InstrumentedModel
from config import Config
from estimators import get_estimator
from models import get_instrumented_model

SEED_SAMPLING = 1
SEED_RANDOM_DIRS = 2
SEED_LINREG = 3
SEED_VISUALIZATION = 5

B = 20
n_clusters = 500

def get_random_dirs(components, dimensions):
    gen = np.random.RandomState(seed=SEED_RANDOM_DIRS)
    dirs = gen.normal(size=(components, dimensions))
    dirs /= np.sqrt(np.sum(dirs**2, axis=1, keepdims=True))
    return dirs.astype(np.float32)

# Compute maximum batch size for given VRAM and network
def get_max_batch_size(inst, device, layer_name=None):
    inst.remove_edits()

    # Reset statistics
    torch.cuda.reset_max_memory_cached(device)
    torch.cuda.reset_max_memory_allocated(device)
    total_mem = torch.cuda.get_device_properties(device).total_memory

    B_max = 20

    # Measure actual usage
    for i in range(2, B_max, 2):
        z = inst.model.sample_latent(n_samples=i)
        if layer_name:
            inst.model.partial_forward(z, layer_name)
        else:
            inst.model.forward(z)

        maxmem = torch.cuda.max_memory_allocated(device)
        del z

        if maxmem > 0.5*total_mem:
            print('Batch size {:d}: memory usage {:.0f}MB'.format(i, maxmem / 1e6))
            return i

    return B_max

# Solve for directions in latent space that match PCs in activaiton space
def linreg_lstsq(comp_np, mean_np, stdev_np, inst, config):
    print('Performing least squares regression', flush=True)

    torch.manual_seed(SEED_LINREG)
    np.random.seed(SEED_LINREG)

    comp = torch.from_numpy(comp_np).float().to(inst.model.device)
    mean = torch.from_numpy(mean_np).float().to(inst.model.device)
    stdev = torch.from_numpy(stdev_np).float().to(inst.model.device)

    n_samp = max(10_000, config.n) // B * B # make divisible
    n_comp = comp.shape[0]
    latent_dims = inst.model.get_latent_dims()
    
    # We're looking for M s.t. M*P*G'(Z) = Z => M*A = Z
    #   Z = batch of latent vectors (n_samples x latent_dims)
    #   G'(Z) = batch of activations at intermediate layer
    #   A = P*G'(Z) = projected activations (n_samples x pca_coords)
    #   M = linear mapping (pca_coords x latent_dims)

    # Minimization min_M ||MA - Z||_l2 rewritten as min_M.T ||A.T*M.T - Z.T||_l2
    # to match format expected by pytorch.lstsq

    # TODO: regression on pixel-space outputs? (using nonlinear optimizer)
    # min_M lpips(G_full(MA), G_full(Z))

    # Tensors to fill with data
    # Dimensions other way around, so these are actually the transposes
    A = np.zeros((n_samp, n_comp), dtype=np.float32)
    Z = np.zeros((n_samp, latent_dims), dtype=np.float32)
    
    # Project tensor X onto PCs, return coordinates
    def project(X, comp):
        N = X.shape[0]
        K = comp.shape[0]
        coords = torch.bmm(comp.expand([N]+[-1]*comp.ndim), X.view(N, -1, 1))
        return coords.reshape(N, K)

    for i in trange(n_samp // B, desc='Collecting samples', ascii=True):
        z = inst.model.sample_latent(B)
        inst.model.partial_forward(z, config.layer)
        act = inst.retained_features()[config.layer].reshape(B, -1)

        # Project onto basis
        act = act - mean
        coords = project(act, comp)
        coords_scaled = coords / stdev

        A[i*B:(i+1)*B] = coords_scaled.detach().cpu().numpy()
        Z[i*B:(i+1)*B] = z.detach().cpu().numpy().reshape(B, -1)

    # Solve least squares fit

    # gelsd = divide-and-conquer SVD; good default
    # gelsy = complete orthogonal factorization; sometimes faster
    # gelss = SVD; slow but less memory hungry
    M_t = scipy.linalg.lstsq(A, Z, lapack_driver='gelsd')[0] # torch.lstsq(Z, A)[0][:n_comp, :]
    
    # Solution given by rows of M_t
    Z_comp = M_t[:n_comp, :]
    Z_mean = np.mean(Z, axis=0, keepdims=True)

    return Z_comp, Z_mean

def regression(comp, mean, stdev, inst, config):
    # Sanity check: verify orthonormality
    M = np.dot(comp, comp.T)
    if not np.allclose(M, np.identity(M.shape[0])):
        det = np.linalg.det(M)
        print(f'WARNING: Computed basis is not orthonormal (determinant={det})')

    return linreg_lstsq(comp, mean, stdev, inst, config)

def compute(config, dump_name, instrumented_model):
    global B

    timestamp = lambda : datetime.datetime.now().strftime("%d.%m %H:%M")
    print(f'[{timestamp()}] Computing', dump_name.name)

    # Ensure reproducibility
    torch.manual_seed(0) # also sets cuda seeds
    np.random.seed(0)

    # Speed up backend
    torch.backends.cudnn.benchmark = True

    has_gpu = torch.cuda.is_available()
    device = torch.device('cuda' if has_gpu else 'cpu')
    layer_key = config.layer

    if instrumented_model is None:
        inst = get_instrumented_model(config.model, config.output_class, layer_key, device)
        model = inst.model
    else:
        print('Reusing InstrumentedModel instance')
        inst = instrumented_model
        model = inst.model
        inst.remove_edits()
        model.set_output_class(config.output_class)

    # Regress back to w space
    if config.use_w:
        print('Using W latent space')
        model.use_w()

    inst.retain_layer(layer_key)
    model.partial_forward(model.sample_latent(1), layer_key)
    sample_shape = inst.retained_features()[layer_key].shape
    sample_dims = np.prod(sample_shape)
    print('Feature shape:', sample_shape)

    input_shape = inst.model.get_latent_shape()
    input_dims = inst.model.get_latent_dims()

    config.components = min(config.components, sample_dims)
    transformer = get_estimator(config.estimator, config.components, config.sparsity)

    X = None
    X_global_mean = None

    # Figure out batch size if not provided
    B = config.batch_size or get_max_batch_size(inst, device, layer_key)

    # Divisible by B (ignored in output name)
    N = config.n // B * B

    # Compute maximum batch size based on RAM + pagefile budget
    target_bytes = 20 * 1_000_000_000 # GB
    feat_size_bytes = sample_dims * np.dtype('float64').itemsize
    N_limit_RAM = np.floor_divide(target_bytes, feat_size_bytes)
    if not transformer.batch_support and N > N_limit_RAM:
        print('WARNING: estimator does not support batching, ' \
            'given config will use {:.1f} GB memory.'.format(feat_size_bytes / 1_000_000_000 * N))

    # 32-bit LAPACK gets very unhappy about huge matrices (in linalg.svd)
    if config.estimator == 'ica':
        lapack_max_N = np.floor_divide(np.iinfo(np.int32).max // 4, sample_dims) # 4x extra buffer
        if N > lapack_max_N:
            raise RuntimeError(f'Matrices too large for ICA, please use N <= {lapack_max_N}')

    print('B={}, N={}, dims={}, N/dims={:.1f}'.format(B, N, sample_dims, N/sample_dims), flush=True)

    # Must not depend on chosen batch size (reproducibility)
    NB = max(B, max(2_000, 3*config.components)) # ipca: as large as possible!
    
    samples = None
    if not transformer.batch_support:
        samples = np.zeros((N + NB, sample_dims), dtype=np.float32)

    torch.manual_seed(config.seed or SEED_SAMPLING)
    np.random.seed(config.seed or SEED_SAMPLING)

    # Use exactly the same latents regardless of batch size
    # Store in main memory, since N might be huge (1M+)
    # Run in batches, since sample_latent() might perform Z -> W mapping
    n_lat = ((N + NB - 1) // B + 1) * B
    latents = np.zeros((n_lat, *input_shape[1:]), dtype=np.float32)
    with torch.no_grad():
        for i in trange(n_lat // B, desc='Sampling latents'):
            latents[i*B:(i+1)*B] = model.sample_latent(n_samples=B).cpu().numpy()

    # Decomposition on non-Gaussian latent space
    samples_are_latents = layer_key in ['g_mapping', 'style'] and inst.model.latent_space_name() == 'W'

    canceled = False
    try:
        X = np.ones((NB, sample_dims), dtype=np.float32)
        action = 'Fitting' if transformer.batch_support else 'Collecting'
        for gi in trange(0, N, NB, desc=f'{action} batches (NB={NB})', ascii=True):
            for mb in range(0, NB, B):
                z = torch.from_numpy(latents[gi+mb:gi+mb+B]).to(device)
                
                if samples_are_latents:
                    # Decomposition on latents directly (e.g. StyleGAN W)
                    batch = z.reshape((B, -1))
                else:
                    # Decomposition on intermediate layer
                    with torch.no_grad():
                        model.partial_forward(z, layer_key)
                    
                    # Permuted to place PCA dimensions last
                    batch = inst.retained_features()[layer_key].reshape((B, -1))

                space_left = min(B, NB - mb)
                X[mb:mb+space_left] = batch.cpu().numpy()[:space_left]

            if transformer.batch_support:
                if not transformer.fit_partial(X.reshape(-1, sample_dims)):
                    break
            else:
                samples[gi:gi+NB, :] = X.copy()
    except KeyboardInterrupt:
        if not transformer.batch_support:
            sys.exit(1) # no progress yet
        
        dump_name = dump_name.parent / dump_name.name.replace(f'n{N}', f'n{gi}')
        print(f'Saving current state to "{dump_name.name}" before exiting')
        canceled = True
        
    if not transformer.batch_support:
        X = samples # Use all samples
        X_global_mean = X.mean(axis=0, keepdims=True, dtype=np.float32) # TODO: activations surely multi-modal...!
        X -= X_global_mean
        
        print(f'[{timestamp()}] Fitting whole batch')
        t_start_fit = datetime.datetime.now()

        transformer.fit(X)
        
        print(f'[{timestamp()}] Done in {datetime.datetime.now() - t_start_fit}')
        assert np.all(transformer.transformer.mean_ < 1e-3), 'Mean of normalized data should be zero'
    else:
        X_global_mean = transformer.transformer.mean_.reshape((1, sample_dims))
        X = X.reshape(-1, sample_dims)
        X -= X_global_mean

    X_comp, X_stdev, X_var_ratio = transformer.get_components()
    
    assert X_comp.shape[1] == sample_dims \
        and X_comp.shape[0] == config.components \
        and X_global_mean.shape[1] == sample_dims \
        and X_stdev.shape[0] == config.components, 'Invalid shape'

    # 'Activations' are really latents in a secondary latent space
    if samples_are_latents:
        Z_comp = X_comp
        Z_global_mean = X_global_mean
    else:
        Z_comp, Z_global_mean = regression(X_comp, X_global_mean, X_stdev, inst, config)

    # Normalize
    Z_comp /= np.linalg.norm(Z_comp, axis=-1, keepdims=True)

    # Random projections
    # We expect these to explain much less of the variance
    random_dirs = get_random_dirs(config.components, np.prod(sample_shape))
    n_rand_samples = min(5000, X.shape[0])
    X_view = X[:n_rand_samples, :].T
    assert np.shares_memory(X_view, X), "Error: slice produced copy"
    X_stdev_random = np.dot(random_dirs, X_view).std(axis=1)

    # Inflate back to proper shapes (for easier broadcasting)
    X_comp = X_comp.reshape(-1, *sample_shape)
    X_global_mean = X_global_mean.reshape(sample_shape)
    Z_comp = Z_comp.reshape(-1, *input_shape)
    Z_global_mean = Z_global_mean.reshape(input_shape)

    # Compute stdev in latent space if non-Gaussian
    lat_stdev = np.ones_like(X_stdev)
    if config.use_w:
        samples = model.sample_latent(5000).reshape(5000, input_dims).detach().cpu().numpy()
        coords = np.dot(Z_comp.reshape(-1, input_dims), samples.T)
        lat_stdev = coords.std(axis=1)

    os.makedirs(dump_name.parent, exist_ok=True)
    np.savez_compressed(dump_name, **{
        'act_comp': X_comp.astype(np.float32),
        'act_mean': X_global_mean.astype(np.float32),
        'act_stdev': X_stdev.astype(np.float32),
        'lat_comp': Z_comp.astype(np.float32),
        'lat_mean': Z_global_mean.astype(np.float32),
        'lat_stdev': lat_stdev.astype(np.float32),
        'var_ratio': X_var_ratio.astype(np.float32),
        'random_stdevs': X_stdev_random.astype(np.float32),
    })

    if canceled:
        sys.exit(1)

    # Don't shutdown if passed as param
    if instrumented_model is None:
        inst.close()
        del inst
        del model

    del X
    del X_comp
    del random_dirs
    del batch
    del samples
    del latents
    torch.cuda.empty_cache()

# Return cached results or commpute if needed
# Pass existing InstrumentedModel instance to reuse it
def get_or_compute(config, model=None, submit_config=None, force_recompute=False):
    if submit_config is None:
        wrkdir = str(Path(__file__).parent.resolve())
        submit_config = SimpleNamespace(run_dir_root = wrkdir, run_dir = wrkdir)
    
    # Called directly by run.py
    return _compute(submit_config, config, model, force_recompute)

def _compute(submit_config, config, model=None, force_recompute=False):
    basedir = Path(submit_config.run_dir)
    outdir = basedir / 'out'
    
    if config.n is None:
        raise RuntimeError('Must specify number of samples with -n=XXX')

    if model and not isinstance(model, InstrumentedModel):
        raise RuntimeError('Passed model has to be wrapped in "InstrumentedModel"')
    
    if config.use_w and not 'StyleGAN' in config.model:
        raise RuntimeError(f'Cannot change latent space of non-StyleGAN model {config.model}')

    transformer = get_estimator(config.estimator, config.components, config.sparsity)
    dump_name = "{}-{}_{}_{}_n{}{}{}.npz".format(
        config.model.lower(),
        config.output_class.replace(' ', '_'),
        config.layer.lower(),
        transformer.get_param_str(),
        config.n,
        '_w' if config.use_w else '',
        f'_seed{config.seed}' if config.seed else ''
    )

    dump_path = basedir / 'cache' / 'components' / dump_name

    if not dump_path.is_file() or force_recompute:
        print('Not cached')
        t_start = datetime.datetime.now()
        compute(config, dump_path, model)
        print('Total time:', datetime.datetime.now() - t_start)
    
    return dump_path