def nadbka(encoder: nn.Module): sum_loss = 0 for i, batch in enumerate(LazyLoader.w300().test_loader): data = batch['data'].to(device) landmarks = batch["meta"]["keypts_normalized"].cuda() content = heatmap_to_measure(encoder(data))[0] eye_dist = landmarks[:, 45] - landmarks[:, 36] eye_dist = eye_dist.pow(2).sum(dim=1).sqrt() sum_loss += ((content - landmarks).pow(2).sum(dim=2).sqrt().mean(dim=1) / eye_dist).sum().item() return sum_loss / len(LazyLoader.w300().test_dataset)
def test(): sum_loss = 0 for i, batch in enumerate(LazyLoader.w300().test_loader): data = batch['data'].to(device) landmarks = batch["meta"]["keypts_normalized"].cuda() content = encoder_HG(data)["coords"] eye_dist = landmarks[:, 45] - landmarks[:, 36] eye_dist = eye_dist.pow(2).sum(dim=1).sqrt() sum_loss += ( (content - landmarks).pow(2).sum(dim=2).sqrt().mean(dim=1) / eye_dist).sum().item() return sum_loss / len(LazyLoader.w300().test_dataset)
def liuboff(encoder: nn.Module): sum_loss = 0 for i, batch in enumerate(LazyLoader.w300().test_loader): data = batch['data'].to(device) landmarks = batch["meta"]["keypts_normalized"].cuda() landmarks[landmarks > 1] = 0.99999 # content = heatmap_to_measure(encoder(data))[0] pred_measure = UniformMeasure2D01(encoder(data)["coords"]) target = UniformMeasure2D01(torch.clamp(landmarks, max=1)) eye_dist = landmarks[:, 45] - landmarks[:, 36] eye_dist = eye_dist.pow(2).sum(dim=1).sqrt() sum_loss += (handmadew1(pred_measure, target) / eye_dist).sum().item() return sum_loss / len(LazyLoader.w300().test_dataset)
def verka_300w_w2(enc): sum_loss = 0 n = len(LazyLoader.w300().test_dataset) for i, batch in enumerate(LazyLoader.w300().test_loader): data = batch['data'].cuda() landmarks = batch["meta"]["keypts_normalized"].cuda() pred = enc(data)["mes"].coord eye_dist = landmarks[:, 45] - landmarks[:, 36] eye_dist = eye_dist.pow(2).sum(dim=1).sqrt() sum_loss += (OTWasDist().forward(pred, landmarks) / eye_dist).sum().item() print("test brule_loss: ", sum_loss / n) return sum_loss / n
def test(enc): sum_loss = 0 for i, batch in enumerate(LazyLoader.w300().test_loader): data = batch['data'].to(device) mes = ProbabilityMeasureFabric(256).from_coord_tensor( batch["meta"]["keypts_normalized"]).cuda() landmarks = batch["meta"]["keypts_normalized"].cuda() content = enc(data) content_xy, _ = heatmap_to_measure(content) eye_dist = landmarks[:, 45] - landmarks[:, 36] eye_dist = eye_dist.pow(2).sum(dim=1).sqrt() sum_loss += ( (content_xy - mes.coord).pow(2).sum(dim=2).sqrt().mean(dim=1) / eye_dist).sum().item() print("test loss: ", sum_loss / len(LazyLoader.w300().test_dataset)) return sum_loss / len(LazyLoader.w300().test_dataset)
def verka_300w_w2_boot(enc): sum_loss = 0 n = len(LazyLoader.w300().test_dataset) loader = torch_data.DataLoader(LazyLoader.w300().test_dataset, batch_size=16, drop_last=False, sampler=torch_data.RandomSampler( LazyLoader.w300().test_dataset, replacement=True, num_samples=n), num_workers=20) for i, batch in enumerate(loader): data = batch['data'].cuda() landmarks = batch["meta"]["keypts_normalized"].cuda() pred = enc(data)["mes"].coord eye_dist = landmarks[:, 45] - landmarks[:, 36] eye_dist = eye_dist.pow(2).sum(dim=1).sqrt() sum_loss += (OTWasDist().forward(pred, landmarks) / eye_dist).sum().item() # print("test brule_loss: ", sum_loss / n) return sum_loss / n
def train(generator, decoder, discriminator, encoder_HG, style_encoder, device, starting_model_number): latent_size = 512 batch_size = 8 sample_z = torch.randn(8, latent_size, device=device) Celeba.batch_size = batch_size W300DatasetLoader.batch_size = batch_size W300DatasetLoader.test_batch_size = 16 test_img = next(LazyLoader.w300().loader_train_inf)["data"][:8].cuda() model = CondStyleGanModel(generator, StyleGANLoss(discriminator), (0.001 / 4, 0.0015 / 4)) style_opt = optim.Adam(style_encoder.parameters(), lr=5e-4, betas=(0.9, 0.99)) cont_opt = optim.Adam(encoder_HG.parameters(), lr=3e-5, betas=(0.5, 0.97)) g_transforms: albumentations.DualTransform = albumentations.Compose([ ToNumpy(), NumpyBatch( albumentations.Compose([ albumentations.ElasticTransform(p=0.7, alpha=150, alpha_affine=1, sigma=10), albumentations.ShiftScaleRotate(p=0.9, rotate_limit=15), ])), ToTensor(device), ]) g_transforms_without_norm: albumentations.DualTransform = albumentations.Compose( [ ToNumpy(), NumpyBatch( albumentations.Compose([ albumentations.ElasticTransform(p=0.3, alpha=150, alpha_affine=1, sigma=10), albumentations.ShiftScaleRotate(p=0.7, rotate_limit=15), ])), ToTensor(device), ]) R_t = DualTransformRegularizer.__call__( g_transforms, lambda trans_dict, img: coord_hm_loss( encoder_HG(trans_dict['image'])["coords"], trans_dict['mask'])) R_s = UnoTransformRegularizer.__call__( g_transforms, lambda trans_dict, img, ltnt: WR.L1("R_s") (ltnt, style_encoder(trans_dict['image']))) barycenter: UniformMeasure2D01 = UniformMeasure2DFactory.load( f"{Paths.default.models()}/face_barycenter_68").cuda().batch_repeat( batch_size) R_b = BarycenterRegularizer.__call__(barycenter, 1.0, 2.0, 4.0) tuner = GoldTuner([0.37, 2.78, 0.58, 1.43, 3.23], device=device, rule_eps=0.001, radius=0.3, active=False) trainer_gan = gan_trainer(model, generator, decoder, encoder_HG, style_encoder, R_s, style_opt, g_transforms) content_trainer = content_trainer_with_gan(cont_opt, tuner, encoder_HG, R_b, R_t, model, generator, g_transforms, decoder, style_encoder) supervise_trainer = content_trainer_supervised( cont_opt, encoder_HG, LazyLoader.w300().loader_train_inf) for i in range(11000): WR.counter.update(i) requires_grad(encoder_HG, False) real_img = next(LazyLoader.celeba().loader).to(device) encoded = encoder_HG(real_img) internal_content = encoded["skeleton"].detach() trainer_gan(i, real_img, internal_content) # content_trainer(real_img) train_content(cont_opt, R_b, R_t, real_img, model, encoder_HG, decoder, generator, style_encoder) supervise_trainer() encoder_ema.accumulate(encoder_HG.module, i, 0.98) if i % 50 == 0 and i > 0: encoder_ema.write_to(encoder_HG.module) if i % 100 == 0: coefs = json.load(open("../parameters/content_loss.json")) print(i, coefs) with torch.no_grad(): # pred_measures_test, sparse_hm_test = encoder_HG(test_img) encoded_test = encoder_HG(test_img) pred_measures_test: UniformMeasure2D01 = UniformMeasure2D01( encoded_test["coords"]) heatmaper_256 = ToGaussHeatMap(256, 1.0) sparse_hm_test_1 = heatmaper_256.forward( pred_measures_test.coord) latent_test = style_encoder(test_img) sparce_mask = sparse_hm_test_1.sum(dim=1, keepdim=True) sparce_mask[sparce_mask < 0.0003] = 0 iwm = imgs_with_mask(test_img, sparce_mask) send_images_to_tensorboard(WR.writer, iwm, "REAL", i) fake_img, _ = generator(encoded_test["skeleton"], [sample_z]) iwm = imgs_with_mask(fake_img, pred_measures_test.toImage(256)) send_images_to_tensorboard(WR.writer, iwm, "FAKE", i) restored = decoder(encoded_test["skeleton"], latent_test) iwm = imgs_with_mask(restored, pred_measures_test.toImage(256)) send_images_to_tensorboard(WR.writer, iwm, "RESTORED", i) content_test_256 = (encoded_test["skeleton"]).repeat(1, 3, 1, 1) * \ torch.tensor([1.0, 1.0, 0.0], device=device).view(1, 3, 1, 1) content_test_256 = (content_test_256 - content_test_256.min() ) / content_test_256.max() send_images_to_tensorboard(WR.writer, content_test_256, "HM", i, normalize=False, range=(0, 1)) if i % 50 == 0 and i >= 0: test_loss = liuboff(encoder_HG) print("liuboff", test_loss) # test_loss = nadbka(encoder_HG) tuner.update(test_loss) WR.writer.add_scalar("liuboff", test_loss, i) if i % 10000 == 0 and i > 0: torch.save( { 'g': generator.module.state_dict(), 'd': discriminator.module.state_dict(), 'c': encoder_HG.module.state_dict(), "s": style_encoder.state_dict(), "e": encoder_ema.storage_model.state_dict() }, f'{Paths.default.models()}/stylegan2_new_{str(i + starting_model_number).zfill(6)}.pt', )
hg = HG_heatmap(heatmapper, num_blocks=1) hg.load_state_dict(weights['gh']) hg = hg.cuda() hm_discriminator = Discriminator(image_size, input_nc=1, channel_multiplier=1) hm_discriminator.load_state_dict(weights["dh"]) hm_discriminator = hm_discriminator.cuda() gan_model_tuda = StyleGanModel[HeatmapToImage](enc_dec.generator, StyleGANLoss(discriminator_img), (0.001/4, 0.0015/4)) gan_model_obratno = StyleGanModel[HG_skeleton](hg, StyleGANLoss(hm_discriminator), (2e-5, 0.0015/4)) style_opt = optim.Adam(enc_dec.style_encoder.parameters(), lr=1e-5) writer = SummaryWriter(f"{Paths.default.board()}/hm2img{int(time.time())}") WR.writer = writer batch = next(LazyLoader.w300().loader_train_inf) test_img = batch["data"].cuda() test_landmarks = torch.clamp(next(LazyLoader.w300_landmarks(args.data_path).loader_train_inf).cuda(), max=1) test_hm = heatmapper.forward(test_landmarks).sum(1, keepdim=True).detach() test_noise = mixing_noise(batch_size, 512, 0.9, device) psp_loss = PSPLoss().cuda() mes_loss = MesBceWasLoss(heatmapper, bce_coef=10000/2, was_coef=100) image_accumulator = Accumulator(enc_dec.generator, decay=0.99, write_every=100) hm_accumulator = Accumulator(hg, decay=0.99, write_every=100) # # fake, _ = enc_dec.generate(test_hm, test_noise) # plt_img = torch.cat([test_img[:3], fake[:3]]).detach().cpu() # plt_lm = torch.cat([hg.forward(test_img)["mes"].coord[:3], test_landmarks[:3]]).detach().cpu()
hg = HG_heatmap(heatmapper) hg.load_state_dict(weights['gh']) hg = hg.cuda() hm_discriminator = Discriminator(image_size, input_nc=1, channel_multiplier=1) hm_discriminator.load_state_dict(weights["dh"]) hm_discriminator = hm_discriminator.cuda() gan_model_tuda = StyleGanModel[HeatmapToImage](enc_dec.generator, StyleGANLoss(discriminator_img), (0.001/4, 0.0015/4)) gan_model_obratno = StyleGanModel[HG_skeleton](hg, StyleGANLoss(hm_discriminator), (2e-5, 0.0015/4)) style_opt = optim.Adam(enc_dec.style_encoder.parameters(), lr=1e-5) writer = SummaryWriter(f"{Paths.default.board()}/hm2img{int(time.time())}") WR.writer = writer test_img = next(LazyLoader.w300().loader_train_inf)["data"].cuda() test_landmarks = torch.clamp(next(LazyLoader.w300augment_landmarks(args.data_path).loader_train_inf).cuda(), max=1) test_hm = heatmapper.forward(test_landmarks).sum(1, keepdim=True).detach() test_noise = mixing_noise(batch_size, 512, 0.9, device) psp_loss = PSPLoss().cuda() mes_loss = MesBceWasLoss(heatmapper, bce_coef=1000000, was_coef=2000) image_accumulator = Accumulator(enc_dec.generator, decay=0.99, write_every=100) hm_accumulator = Accumulator(hg, decay=0.99, write_every=100) for i in range(100000): WR.counter.update(i)
ToTensor(device) ]) R_s = UnoTransformRegularizer.__call__( g_transforms, lambda trans_dict, img, ltnt: Loss(nn.L1Loss()( ltnt, style_encoder(trans_dict['image'])))) W300DatasetLoader.batch_size = 24 W300DatasetLoader.test_batch_size = 64 Celeba.batch_size = 24 heatmaper = ToHeatMap(64) # tuner = GoldTuner([1.0, 1.0], device=device, rule_eps=0.02, radius=0.5, active=True) w300_test = next(iter(LazyLoader.w300().test_loader)) w300_test_image = w300_test['data'].to(device)[:8] def hm_svoego_roda_loss(pred, target, coef=1.0, l1_coef=0.0): pred_mes = UniformMeasure2DFactory.from_heatmap(pred) target_mes = UniformMeasure2DFactory.from_heatmap(target) # pred = pred.relu() + 1e-15 # target[target < 1e-7] = 0 # target[target > 1 - 1e-7] = 1 if torch.isnan(pred).any() or torch.isnan(target).any(): print("nan in hm") return Loss.ZERO()
# style_encoder = style_encoder.cuda() decoder = CondGenDecode(generator) # generator = nn.DataParallel(generator, [0, 1, 3]) # discriminator = nn.DataParallel(discriminator, [0, 1, 3]) # encoder_HG = nn.DataParallel(encoder_HG, [0, 1, 3]) # decoder = nn.DataParallel(decoder, [0, 1, 3]) test_img = next(LazyLoader.celeba().loader)[:4].cuda() heatmaper = ToGaussHeatMap(64, 1.5) sample_z = torch.randn(4, 512, device=device) noise = mixing_noise(4, 512, 0.9, device) LazyLoader.w300_save = None W300DatasetLoader.test_batch_size = 4 w300_test = next(iter(LazyLoader.w300().test_loader)) w300_test_image = w300_test['data'].to(device)[:4] w300_test_mask = ProbabilityMeasureFabric(256).from_coord_tensor( w300_test["meta"]["keypts_normalized"][:4].to(device)) sparse_hm = heatmaper.forward(w300_test_mask.coord * 63).detach() g_transforms: albumentations.DualTransform = albumentations.Compose([ ToNumpy(), NumpyBatch( albumentations.Compose([ ResizeMask(h=256, w=256), albumentations.ElasticTransform(p=1, alpha=150, alpha_affine=1, sigma=10), # albumentations.ShiftScaleRotate(p=0.7, rotate_limit=15),
from torchvision.utils import make_grid from dataset.lazy_loader import LazyLoader, W300DatasetLoader, CelebaWithKeyPoints, Celeba from sklearn.neighbors import NearestNeighbors import numpy as np import matplotlib.pyplot as plt from dataset.toheatmap import ToGaussHeatMap from dataset.probmeasure import UniformMeasure2D01 import pandas as pd import networkx as nx import ot from barycenters.sampler import Uniform2DBarycenterSampler, Uniform2DAverageSampler, ImageBarycenterSampler from parameters.path import Paths from joblib import Parallel, delayed N = 100 dataset = LazyLoader.w300().dataset_train D = np.load(f"{Paths.default.models()}/w300graph{N}.npy") padding = 68 prob = np.ones(padding) / padding NS = 7000 def LS(k): return dataset[k]["meta"]['keypts_normalized'].numpy() ls = [] images = [] for k in range(N): dk = dataset[k]
def hm_svoego_roda_loss(pred, target): pred_xy, _ = heatmap_to_measure(pred) t_xy, _ = heatmap_to_measure(target) return Loss(nn.BCELoss()(pred, target) + nn.MSELoss()(pred_xy, t_xy) * 0.0005 + (pred - target).abs().mean() * 0.3) R_t = DualTransformRegularizer.__call__( g_transforms, lambda trans_dict, img: hm_svoego_roda_loss( encoder_HG(trans_dict['image']), trans_dict['mask'])) for epoch in range(130): for i, batch in enumerate(LazyLoader.w300().loader_train): # print(i) counter.update(i + epoch * len(LazyLoader.w300().loader_train)) data = batch['data'].to(device) mes = ProbabilityMeasureFabric(256).from_coord_tensor( batch["meta"]["keypts_normalized"]).cuda() target_hm = heatmaper.forward(mes.probability, mes.coord * 63) content = encoder_HG(data) hm_svoego_roda_loss(content, target_hm).minimize_step(cont_opt) if i % 5 == 0: real_img = next(LazyLoader.celeba().loader).to(device) content = encoder_HG(data) coefs = json.load(open("../parameters/content_loss_sup.json"))
eye_dist = landmarks[:, 45] - landmarks[:, 36] eye_dist = eye_dist.pow(2).sum(dim=1).sqrt() sum_loss += ( (content - landmarks).pow(2).sum(dim=2).sqrt().mean(dim=1) / eye_dist).sum().item() return sum_loss / len(LazyLoader.w300().test_dataset) cont_opt = optim.Adam(encoder_HG.parameters(), lr=3e-5, betas=(0.5, 0.97)) W300DatasetLoader.batch_size = 16 W300DatasetLoader.test_batch_size = 32 supervise_trainer = content_trainer_supervised( cont_opt, encoder_HG, LazyLoader.w300().loader_train_inf) test_img = next(LazyLoader.w300().loader_train_inf)['data'].cuda() for i in range(10000): supervise_trainer() if i % 100 == 0: with torch.no_grad(): liuboff = test() print(liuboff) WR.writer.add_scalar("liuboff", liuboff, i) encoded_test = encoder_HG(test_img) pred_measures_test: UniformMeasure2D01 = UniformMeasure2D01(