Guccio-AI-Designer / decomposition.py
2Much2Code:)
commit
ba5dcdc
# 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