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: