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,