from time import time from os.path import join import os import sys import lpips from GAN_hessian_compute import hessian_compute, get_full_hessian from torchvision.transforms import ToPILImage from torchvision.utils import make_grid from GAN_utils import loadBigGAN, loadStyleGAN2, BigGAN_wrapper, StyleGAN2_wrapper, loadPGGAN, PGGAN_wrapper from hessian_analysis_tools import plot_spectra, compute_hess_corr from lpips import LPIPS ImDist = LPIPS(net="squeeze") datadir = r"E:\OneDrive - Washington University in St. Louis\HessNetArchit\PGGAN" os.makedirs(datadir, exist_ok=True) #%% PGGAN = loadPGGAN() SD = PGGAN.state_dict() #%% shuffled_SD = {} for name, Weight in SD.items(): idx = torch.randperm(Weight.numel()) W_shuf = Weight.view(-1)[idx].view(Weight.shape) shuffled_SD[name] = W_shuf #%% torch.save(shuffled_SD, join(datadir, "PGGAN_shuffle.pt")) # print(name, Weight.shape, Weight.mean().item(), Weight.std().item()) #%% PGGAN_sf = loadPGGAN() PGGAN_sf.load_state_dict(torch.load(join(datadir, "PGGAN_shuffle.pt"))) G_sf = PGGAN_wrapper(PGGAN_sf) #%%
from load_hessian_data import load_Haverage import matplotlib.pylab as plt import torch, numpy as np from os.path import join import os from torchvision.utils import make_grid from torchvision.transforms import ToPILImage import matplotlib.pylab as plt from torch_utils import show_imgrid, save_imgrid saveroot = r"E:\OneDrive - Washington University in St. Louis\GAN_baddots" #%% Progressive Growing GAN modelsnm = "PGGAN" savedir = join(saveroot, modelsnm) os.makedirs(savedir, exist_ok=True) PGAN = loadPGGAN() G = PGGAN_wrapper(PGAN) H, eva, evc = load_Haverage("PGGAN", descend=True) randvec = G.sample_vector(10, 'cpu') distnorm = randvec.norm(dim=1).mean().item() # distnorm = np.sqrt(randvec.shape[1]) #% baddot_imgs_pos = G.visualize_batch_np(distnorm * evc[:, :20].T) show_imgrid( baddot_imgs_pos, nrow=5, padding=2, ) save_imgrid( baddot_imgs_pos, join(savedir, "GAN_baddots_pos.png"),
figdir=figdir, namestr="SG2_%s_orig" % (modelsnm, ), page_B=10, eig_rng=(0, 60), maxdist=9, rown=5, sphere=False, transpose=False, RND=0) # for eigidx in range(40): # vis_eigen_action(evc_avg[:, -eigidx - 1], ref_codes, SG, figdir=figdir, # namestr="SG2_%s_eig%d_orig" % (modelsnm, eigidx + 1), # maxdist=9, rown=7, sphere=False, transpose=False, RND=0) #%% PGGAN PGGAN = loadPGGAN() # , size=256 PG = PGGAN_wrapper(PGGAN) figdir = join(axesdir, "PGGAN") os.makedirs(figdir, exist_ok=True) data = np.load( join(rootdir, "PGGAN", "H_avg_%s.npz" % "PGGAN" )) #, H_avg=H_avg, eva_avg=eva_avg, evc_avg=evc_avg, feats=feat_col) H_avg, eva_avg, evc_avg = data["H_avg"], data["eva_avg"], data["evc_avg"] #%% RND = np.random.randint(1E4) ref_codes = np.random.randn(8, 512) for eigidx in range(40): vis_eigen_action(evc_avg[:, -eigidx - 1], ref_codes, PG, figdir=figdir,