except Exception as e: print(e) continue coefs = json.load(open(os.path.join(sys.path[0], "../parameters/cycle_loss_2.json"))) fake, fake_latent = enc_dec.generate(heatmap_sum) fake_latent_pred = enc_dec.encode_latent(fake) gan_model_tuda.discriminator_train([real_img], [fake.detach()]) ( gan_model_tuda.generator_loss([real_img], [fake]) + l1_loss(fake_latent_pred, fake_latent) * coefs["style"] ).minimize_step(gan_model_tuda.optimizer.opt_min, style_opt) hm_pred = hg.forward(real_img)["hm_sum"] hm_ref = heatmapper.forward(landmarks).detach().sum(1, keepdim=True) gan_model_obratno.discriminator_train([hm_ref], [hm_pred.detach()]) gan_model_obratno.generator_loss([hm_ref], [hm_pred]).__mul__(coefs["obratno"]) \ .minimize_step(gan_model_obratno.optimizer.opt_min) fake2, _ = enc_dec.generate(heatmap_sum) WR.writable("cycle", mes_loss.forward)(hg.forward(fake2)["mes"], UniformMeasure2D01(landmarks)).__mul__(coefs["hm"]) \ .minimize_step(gan_model_tuda.optimizer.opt_min, gan_model_obratno.optimizer.opt_min) latent = enc_dec.encode_latent(g_transforms(image=real_img)["image"]) restored = enc_dec.decode(hg.forward(real_img)["hm_sum"], latent) WR.writable("cycle2", psp_loss.forward)(real_img, real_img, restored, latent).__mul__(coefs["img"]) \ .minimize_step(gan_model_tuda.optimizer.opt_min, gan_model_obratno.optimizer.opt_min, style_opt) image_accumulator.step(i)
# 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) batch = next(LazyLoader.cardio().loader_train_inf) real_img = batch["image"].cuda() train_landmarks = batch["keypoints"].cuda() coefs = json.load(open(os.path.join(sys.path[0], "../parameters/cycle_loss_2.json"))) WR.writable("sup", mes_loss.forward)(hg.forward(real_img)["mes"], UniformMeasure2D01(train_landmarks)).__mul__(coefs["sup"]) \ .minimize_step(cont_opt) hm_accumulator.step(i) if i % 1000 == 0 and i > 0: torch.save( { 'gh': hg.state_dict(), }, f'{Paths.default.models()}/cardio_brule_sup_{str(i + starting_model_number).zfill(6)}.pt', ) if i % 100 == 0: print(i)
for i in range(100000): WR.counter.update(i) batch = next(LazyLoader.cardio().loader_train_inf) real_img = batch["image"].cuda() train_landmarks = batch["keypoints"].cuda() coefs = json.load( open(os.path.join(sys.path[0], "../parameters/cycle_loss_2.json"))) # WR.writable("sup", mes_loss.forward)(hg.forward(real_img)["mes"], UniformMeasure2D01(train_landmarks)).__mul__(coefs["sup"]) \ # .minimize_step(cont_opt) with torch.no_grad(): pred = hg.forward(real_img) hm_pred = pred["hm_sum"].detach() mes_pred = pred["mes"].detach() fake, fake_latent = enc_dec.generate(hm_pred) fake_latent_pred = enc_dec.encode_latent(fake) gan_model_tuda.discriminator_train([real_img], [fake.detach()]) (gan_model_tuda.generator_loss([real_img], [fake]) + l1_loss(fake_latent_pred, fake_latent) * coefs["style"]).minimize_step( gan_model_tuda.optimizer.opt_min, style_opt) train_content(cont_opt, R_b, R_t, real_img, hg) fake2, _ = enc_dec.generate(hm_pred) WR.writable("cycle", mes_loss.forward)(hg.forward(fake2)["mes"], mes_pred).__mul__(coefs["hm"]) \
batch = next(LazyLoader.human36(use_mask=True).loader_train_inf) real_img = batch["A"].cuda() landmarks = torch.clamp(batch["paired_B"].cuda(), min=0, max=1) heatmap = heatmapper.forward(landmarks).detach() coefs = json.load( open(os.path.join(sys.path[0], "../parameters/cycle_loss.json"))) fake, fake_latent = enc_dec.generate(heatmap) gan_model_tuda.discriminator_train([real_img], [fake.detach()]) (gan_model_tuda.generator_loss([real_img], [fake])).minimize_step( gan_model_tuda.optimizer.opt_min) fake2, _ = enc_dec.generate(heatmap) WR.writable("cycle", mes_loss.forward)(hg.forward(fake2)["mes"], UniformMeasure2D01(landmarks)).__mul__(coefs["hm"]) \ .minimize_step(gan_model_tuda.optimizer.opt_min) image_accumulator.step(i) if i % 10000 == 0 and i > 0: torch.save( { 'gi': enc_dec.generator.state_dict(), 'di': discriminator_img.state_dict(), 's': enc_dec.style_encoder.state_dict() }, f'{Paths.default.models()}/human_gan_{str(i + starting_model_number).zfill(6)}.pt', ) if i % 100 == 0: