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:
