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)
def hm_svoego_roda_loss(pred, target): pred_xy, _ = heatmap_to_measure(pred) with torch.no_grad(): t_xy, _ = heatmap_to_measure(target) return Loss( nn.BCELoss()(pred, target) + nn.MSELoss()(pred_xy, t_xy) * 0.001 # (pred - target).abs().mean() * 0.1 )
def compute(content: Tensor, target_hm: Tensor): content_xy, _ = heatmap_to_measure(content) target_xy, _ = heatmap_to_measure(target_hm) lossyash = Loss(nn.BCELoss()(content, target_hm) * weight + nn.MSELoss()(content_xy, target_xy) * weight * 0.001) if name: writer.add_scalar(name, lossyash.item(), counter.get_iter(name)) return lossyash
def hm_svoego_roda_loss(pred, target): pred_coord = heatmap_to_measure(pred)[0] target_coord = heatmap_to_measure(target)[0] 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(): return Loss.ZERO() bce = nn.BCELoss()(pred, target) if torch.isnan(bce).any(): return Loss.ZERO() return Loss(bce + nn.MSELoss()(pred_coord, target_coord) * 0.0005)
def __getitem__(self, index): keypts = np.load(self.data[index]) # numpy - [200, 2] hm = self.heatmapper.forward(torch.tensor(keypts)[None, ])[0] transformed = self.transform(image=np.zeros_like( np.array(hm.permute(1, 2, 0))), mask=np.array(hm.permute(1, 2, 0))) mask = transformed["mask"] coord, p = heatmap_to_measure(mask.permute(2, 0, 1)[None]) return coord[0].type(torch.float32) #torch.tensor - [200, 2]
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 forward(self, image: Tensor): B, C, D, D = image.shape heatmaps: List[Tensor] = self.model.forward(image) out = heatmaps[-1] hm = self.up(self.postproc(out)) coords, _ = heatmap_to_measure(hm) return { "out": out, "mes": UniformMeasure2D01(coords), 'hm': hm, "softmax": self.postproc(out) }
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 return_coords(self, image: Tensor): heatmaps = self.forward(image) coords, p = heatmap_to_measure(heatmaps) return coords
def from_heatmap(hm: Tensor) -> UniformMeasure2D01: coord, _ = heatmap_to_measure(hm) return UniformMeasure2D01(coord)
def train(generator, discriminator, encoder, style_encoder, device, starting_model_number): batch = 32 Celeba.batch_size = batch latent_size = 512 model = CondStyleGanModel(generator, StyleGANLoss(discriminator), (0.001, 0.0015)) style_opt = optim.Adam(style_encoder.parameters(), lr=5e-4, betas=(0.5, 0.97)) g_transforms: albumentations.DualTransform = albumentations.Compose([ ToNumpy(), NumpyBatch( albumentations.ElasticTransform(p=0.8, alpha=150, alpha_affine=1, sigma=10)), NumpyBatch(albumentations.ShiftScaleRotate(p=0.5, rotate_limit=10)), ToTensor(device) ]) R_s = UnoTransformRegularizer.__call__( g_transforms, lambda trans_dict, img, ltnt: L1("R_s") (ltnt, style_encoder(trans_dict['image']))) sample_z = torch.randn(batch, latent_size, device=device) test_img = next(LazyLoader.celeba().loader).to(device) print(test_img.shape) test_cond = encoder(test_img) requires_grad(encoder, False) # REMOVE BEFORE TRAINING t_start = time.time() for i in range(100000): counter.update(i) real_img = next(LazyLoader.celeba().loader).to(device) img_content = encoder(real_img).detach() noise = mixing_noise(batch, latent_size, 0.9, device) fake, _ = generator(img_content, noise) model.discriminator_train([real_img], [fake.detach()], img_content) writable("Generator loss", model.generator_loss)([real_img], [fake], [], img_content)\ .minimize_step(model.optimizer.opt_min) # print("gen train", time.time() - t1) if i % 5 == 0 and i > 0: noise = mixing_noise(batch, latent_size, 0.9, device) img_content = encoder(real_img).detach() fake, fake_latent = generator(img_content, noise, return_latents=True) fake_latent_test = fake_latent[:, [0, 13], :].detach() fake_latent_pred = style_encoder(fake) fake_content_pred = encoder(fake) restored = generator.module.decode( img_content[:batch // 2], style_encoder(real_img[:batch // 2])) (HMLoss("BCE content gan", 5000)(fake_content_pred, img_content) + L1("L1 restored")(restored, real_img[:batch // 2]) * 50 + L1("L1 style gan")(fake_latent_pred, fake_latent_test) * 30 + R_s(fake.detach(), fake_latent_pred) * 50).minimize_step( model.optimizer.opt_min, style_opt) if i % 100 == 0: t_100 = time.time() print(i, t_100 - t_start) t_start = time.time() with torch.no_grad(): fake_img, _ = generator(test_cond, [sample_z]) coords, p = heatmap_to_measure(test_cond) test_mes = ProbabilityMeasure(p, coords) iwm = imgs_with_mask(fake_img, test_mes.toImage(256)) send_images_to_tensorboard(writer, iwm, "FAKE", i) iwm = imgs_with_mask(test_img, test_mes.toImage(256)) send_images_to_tensorboard(writer, iwm, "REAL", i) restored = generator.module.decode(test_cond, style_encoder(test_img)) send_images_to_tensorboard(writer, restored, "RESTORED", i) if i % 10000 == 0 and i > 0: torch.save( { 'g': generator.state_dict(), 'd': discriminator.state_dict(), 'style': style_encoder.state_dict() # 'enc': cont_style_encoder.state_dict(), }, f'/trinity/home/n.buzun/PycharmProjects/saved/stylegan2_w300_{str(starting_model_number + i).zfill(6)}.pt', )
# HMLoss("BCE content gan", 5000)(fake_content_pred, img_content.detach()) + # Loss(nn.L1Loss()(restored, real_img[:W300DatasetLoader.batch_size//2]) * 50) + # Loss(nn.L1Loss()(fake_latent_pred, fake_latent_test) * 25) + # R_s(fake.detach(), fake_latent_pred) * 50 # ).minimize_step( # model.optimizer.opt_min, # style_opt, # ) # img_content = encoder_HG(real_img) # fake, fake_latent = generator(img_content, noise, return_latents=True) # fake_content_pred = encoder_HG(fake) # # # disc_influence = model.loss.generator_loss(real=None, fake=[real_img, img_content]) * 2 # (HMLoss("BCE content gan", 1)(fake_content_pred, img_content.detach()) + # disc_influence).minimize_step(enc_opt) if i % 50 == 0 and i > 0: with torch.no_grad(): test_loss = test(encoder_HG) print(test_loss) # tuner.update(test_loss) coord, p = heatmap_to_measure(encoder_HG(w300_test_image)) pred_measure = ProbabilityMeasure(p, coord) iwm = imgs_with_mask(w300_test_image, pred_measure.toImage(256)) send_images_to_tensorboard(writer, iwm, "W300_test_image", i) writer.add_scalar("test_loss", test_loss, i) # torch.save(enc.state_dict(), f"/home/ibespalov/pomoika/hg2_e{epoch}.pt")
def train(generator, decoder, discriminator, discriminatorHG, skeleton_encoder, style_encoder, device, starting_model_number): latent_size = 512 batch_size = 12 sample_z = torch.randn(8, latent_size, device=device) Celeba.batch_size = batch_size W300DatasetLoader.batch_size = batch_size W300DatasetLoader.test_batch_size = 16 encoder_HG_supervised = IXHG(num_classes=68, heatmap_size=64) encoder_HG_supervised.load_state_dict(torch.load(f'{Paths.default.models()}/hg2_e29.pt', map_location="cpu")) encoder_HG_supervised = encoder_HG_supervised.cuda() requires_grad(encoder_HG_supervised, False) style_opt = optim.Adam(style_encoder.parameters(), lr=5e-4, betas=(0.9, 0.99)) test_img = next(LazyLoader.celeba().loader)[:8].cuda() loss_st: StyleGANLoss = StyleGANLoss(discriminator) # model = CondStyleGanModel(generator, loss_st, (0.001, 0.0015)) # loss_st2: StyleGANLoss = StyleGANLoss(discriminatorHG) model2 = CondStyleGanModel(skeleton_encoder, loss_st2, (0.0001, 0.001)) skeletoner = CoordToGaussSkeleton(size, 4) # tuda_trainer = gan_tuda_trainer(model, generator) obratno_trainer = gan_obratno_trainer(model2) # tuda_obratno_trainer = gan_tuda_obratno_trainer(generator, encoder_HG, decoder, style_encoder, style_opt) for i in range(100000): counter.update(i) # requires_grad(encoder_HG, False) # REMOVE BEFORE TRAINING real_img = next(LazyLoader.celeba().loader).to(device) heatmap = encoder_HG_supervised.get_heatmaps(real_img).detach() coords, p = heatmap_to_measure(heatmap) skeleton = skeletoner.forward(coords).sum(dim=1, keepdim=True) # tuda_trainer(real_img, heatmap) obratno_trainer(real_img, skeleton) # tuda_obratno_trainer(real_img, heatmap) # if i % 10 == 0: # accumulate(encoder_ema, encoder_HG.module, 0.95) # pred_hm = encoder_HG(real_img) # stariy_hm_loss(pred_hm, heatmap).minimize_step(model2.optimizer.opt_min) if i % 100 == 0 and i >= 0: with torch.no_grad(): content_test = skeleton_encoder(test_img)["skeleton"].sum(dim=1, keepdim=True) iwm = imgs_with_mask(test_img, content_test, border=0.1) send_images_to_tensorboard(writer, iwm, "REAL", i) # fake_img, _ = generator(sparse_hm_test, [sample_z]) # iwm = imgs_with_mask(fake_img, pred_measures_test.toImage(256)) # send_images_to_tensorboard(writer, iwm, "FAKE", i) # # restored = decoder(sparse_hm_test, latent_test) # iwm = imgs_with_mask(restored, pred_measures_test.toImage(256)) # send_images_to_tensorboard(writer, iwm, "RESTORED", i) content_test_hm = (content_test - content_test.min()) / content_test.max() send_images_to_tensorboard(writer, content_test_hm, "HM", i, normalize=False, range=(0, 1)) heatmap_test = encoder_HG_supervised.get_heatmaps(test_img).detach() coords_test, _ = heatmap_to_measure(heatmap_test) skeleton_test = skeletoner.forward(coords_test).sum(dim=1, keepdim=True) iwm = imgs_with_mask(test_img, skeleton_test, border=0.1) send_images_to_tensorboard(writer, iwm, "REF", i) if i % 50 == 0 and i >= 0: test_loss = liuboff(skeleton_encoder) print("liuboff", test_loss) # test_loss = nadbka(encoder_HG) # tuner.update(test_loss) 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': skeleton_encoder.module.state_dict(), "s": style_encoder.state_dict(), "d2": discriminatorHG.module.state_dict(), # "ema": encoder_ema.state_dict() }, f'{Paths.default.models()}/cyclegan_{str(i + starting_model_number).zfill(6)}.pt', )
def forward(self, hm: Tensor): coord, p = heatmap_to_measure(hm) return self.net(coord.view(-1, 136))
barycenter.requires_grad_() coord = barycenter.coord opt = optim.Adam(iter([coord]), lr=0.0006) encoder_HG = HG_softmax2020(num_classes=68, heatmap_size=64) encoder_HG.load_state_dict( torch.load(f"{Paths.default.models()}/hg2_e29.pt", map_location="cpu")) encoder_HG = encoder_HG.cuda() for iter in range(3000): img = next(LazyLoader.celeba().loader).cuda() content = encoder_HG(img) coord, p = heatmap_to_measure(content) mes = ProbabilityMeasure(p, coord) barycenter_cat = fabric.cat([barycenter] * batch_size) loss = Samples_Loss()(barycenter_cat, mes) opt.zero_grad() loss.to_tensor().backward() opt.step() barycenter.probability.data = barycenter.probability.relu().data barycenter.probability.data /= barycenter.probability.sum(dim=1, keepdim=True) if iter % 100 == 0: