def _run_epoch(self, dataset, eval=False):
epoch_starttime = time.time()
batchsize = self.args.batchsize_eval if eval else self.args.batchsize
self.iters_per_epoch = int(len(dataset) / batchsize)
self.iter_starttime = time.time()
self.iter_in_epoch = 0
dataloader = td.DataLoader(
dataset,
batch_size=batchsize,
shuffle=not eval,
num_workers=self.workers,
drop_last=not eval,
)
batch_samples = None
for data in dataloader:
batch = Batch(data, eval=eval)
batch_samples = self._run_batch(batch, eval=eval)
self.total_iter += 1
self.saae.total_iter = self.total_iter
self.iter_in_epoch += 1
self._print_epoch_summary(self.epoch_stats,
epoch_starttime,
eval=eval,
batch_samples=batch_samples)
result['landmarks'], self.landmark_sigma, self.landmark_ids)
return result
if __name__ == '__main__':
from utils.nn import Batch, to_numpy
import utils.common
from landmarks import lmconfig as lmcfg
utils.common.init_random(3)
lmcfg.config_landmarks('wflw')
ds = WFLW(train=True, deterministic=True, use_cache=True, daug=0)
# ds.filter_labels({'pose':0, 'blur':0, 'occlusion':1})
dl = td.DataLoader(ds, batch_size=1, shuffle=False, num_workers=0)
cfg.WITH_LANDMARK_LOSS = False
for data in dl:
batch = Batch(data, gpu=False)
images = vis._to_disp_images(batch.images, denorm=True)
# lms = lmutils.convert_landmarks(to_numpy(batch.landmarks), lmutils.LM98_TO_LM68)
lms = batch.landmarks
images = vis.add_landmarks_to_images(images,
lms,
draw_wireframe=False,
color=(0, 255, 0),
radius=3)
vis.vis_square(images, nCols=1, fx=1., fy=1., normalize=False)
def get_fixed_samples(ds, num):
dl = td.DataLoader(ds, batch_size=num, shuffle=False, num_workers=0)
data = next(iter(dl))
return Batch(data, n=num)
parser.add_argument('--st', default=0, type=int)
parser.add_argument('--nd', default=1, type=int)
args = parser.parse_args()
extract_features(st=args.st, nd=args.nd)
if __name__ == '__main__':
# extract_main()
# exit()
import torch
from utils import vis, face_processing
from utils.nn import Batch
torch.manual_seed(0)
torch.cuda.manual_seed_all(0)
ds = LFWImages(train=True, deterministic=True, max_samples=40, use_cache=False)
dl = torch.utils.data.DataLoader(ds, batch_size=10, shuffle=True, num_workers=1)
data = next(iter(dl))
inputs = Batch(data).images
ds_utils.denormalize(inputs)
# imgs = vis.add_id_to_images(inputs.numpy(), data[1].numpy())
# imgs = vis.add_pose_to_images(inputs.numpy(), pose.numpy())
# imgs = vis.add_landmarks_to_images(inputs.numpy(), landmarks.numpy())
imgs = inputs.detach().cpu().numpy()
vis.vis_square(imgs, fx=1.0, fy=1.0, normalize=False)
args = parser.parse_args()
extract_features(st=args.st, nd=args.nd)
if __name__ == '__main__':
# extract_main()
# create_annotations(split='dev', num_ids=500)
# extract_crops()
# read_openface_csvs()
# exit()
from utils import vis, face_processing
ds = VoxCeleb(train=True, max_samples=50000, use_cache=True)
print(ds)
dl = td.DataLoader(ds, batch_size=40, shuffle=False, num_workers=0)
from utils.nn import Batch
for data in dl:
batch = Batch(data)
print(batch.ids)
ds_utils.denormalize(batch.images)
# vis.show_images_in_batch(batch.images.detach().cpu())
vis.vis_square(batch.images.detach().cpu(),
fx=0.7,
fy=0.7,
normalize=False)
# print(item)
def _run_batch(self, data, train_autoencoder, eval=False, ds=None):
time_dataloading = time.time() - self.iter_starttime
time_proc_start = time.time()
iter_stats = {'time_dataloading': time_dataloading}
batch = Batch(data, eval=eval)
self.saae.zero_grad()
self.saae.eval()
input_images = batch.images_mod if batch.images_mod is not None else batch.images
with torch.set_grad_enabled(self.args.train_encoder):
z_sample = self.saae.Q(input_images)
iter_stats.update({'z_recon_mean': z_sample.mean().item()})
#######################
# Reconstruction phase
#######################
with torch.set_grad_enabled(self.args.train_encoder and not eval):
X_recon = self.saae.P(z_sample)
with torch.no_grad():
diff = torch.abs(batch.images - X_recon) * 255
loss_recon_l1 = torch.mean(diff)
loss_Q = loss_recon_l1 * cfg.W_RECON
iter_stats['loss_recon'] = loss_recon_l1.item()
l1_dist_per_img = diff.reshape(len(batch.images), -1).mean(dim=1)
iter_stats['l1_recon_errors'] = to_numpy(l1_dist_per_img)
#######################
# Landmark predictions
#######################
train_lmhead = not eval and not args.train_coords
lm_preds_max = None
with torch.set_grad_enabled(train_lmhead):
self.saae.LMH.train(train_lmhead)
X_lm_hm = self.saae.LMH(self.saae.P)
if batch.lm_heatmaps is not None:
loss_lms = F.mse_loss(batch.lm_heatmaps, X_lm_hm) * 100 * 3
if (eval or self._is_printout_iter()):
# expensive, so only calculate when every N iterations
X_lm_hm = lmutils.decode_heatmap_blob(X_lm_hm)
X_lm_hm = lmutils.smooth_heatmaps(X_lm_hm)
lm_preds_max = self.saae.heatmaps_to_landmarks(X_lm_hm)
iter_stats.update({'loss_Q': loss_Q.item()})
if not args.train_coords:
iter_stats.update({'loss_lms': loss_lms.item()})
if (eval or self._is_printout_iter()):
lm_gt = to_numpy(batch.landmarks)
lm_errs_max = lmutils.calc_landmark_nme_per_img(lm_gt,
lm_preds_max,
ocular_norm=self.args.ocular_norm,
landmarks_to_eval=lmcfg.LANDMARKS_NO_OUTLINE)
lm_errs_max_outline = lmutils.calc_landmark_nme_per_img(lm_gt,
lm_preds_max,
ocular_norm=self.args.ocular_norm,
landmarks_to_eval=lmcfg.LANDMARKS_ONLY_OUTLINE)
lm_errs_max_all = lmutils.calc_landmark_nme_per_img(lm_gt,
lm_preds_max,
ocular_norm=self.args.ocular_norm,
landmarks_to_eval=lmcfg.ALL_LANDMARKS)
nmes = lmutils.calc_landmark_nme(lm_gt, lm_preds_max, ocular_norm=self.args.ocular_norm)
# nccs = lmutils.calc_landmark_ncc(batch.images, X_recon, lm_gt)
iter_stats.update({'lm_errs_max': lm_errs_max,
'lm_errs_max_all': lm_errs_max_all,
'lm_errs_max_outline': lm_errs_max_outline,
'nmes': nmes,
# 'nccs': nccs
})
if train_lmhead:
if self.args.train_encoder:
loss_lms = loss_lms * 80.0
loss_lms.backward()
self.optimizer_lm_head.step()
if self.args.train_encoder:
self.optimizer_Q.step()
# statistics
iter_stats.update({'epoch': self.epoch, 'timestamp': time.time(),
'iter_time': time.time() - self.iter_starttime,
'time_processing': time.time() - time_proc_start,
'iter': self.iter_in_epoch, 'total_iter': self.total_iter, 'batch_size': len(batch)})
self.iter_starttime = time.time()
self.epoch_stats.append(iter_stats)
# print stats every N mini-batches
if self._is_printout_iter():
self._print_iter_stats(self.epoch_stats[-self.print_interval:])
# Batch visualization
#
if self._is_printout_iter():
f = 2.0 / cfg.INPUT_SCALE_FACTOR
# lmutils.visualize_random_faces(self.saae, 20, 0)
lmutils.visualize_batch(batch.images, batch.landmarks, X_recon, X_lm_hm, lm_preds_max,
lm_heatmaps=batch.lm_heatmaps,
images_mod=batch.images_mod,
ds=ds, wait=self.wait,
landmarks_to_draw=lmcfg.ALL_LANDMARKS,
ocular_norm=args.ocular_norm,
f=f,
clean=False,
overlay_heatmaps_input=False,
overlay_heatmaps_recon=False)
def _run_batch(self, data, eval=False, ds=None):
time_dataloading = time.time() - self.iter_starttime
time_proc_start = time.time()
iter_stats = {"time_dataloading": time_dataloading}
batch = Batch(data, eval=eval)
X_target = (batch.target_images
if batch.target_images is not None else batch.images)
self.saae.zero_grad()
loss = torch.zeros(1, requires_grad=True).cuda()
#######################
# Encoding
#######################
with torch.set_grad_enabled(self.args.train_encoder):
z_sample = self.saae.Q(batch.images)
###########################
# Encoding regularization
###########################
if ((not eval or self._is_printout_iter(eval))
and self.args.with_zgan and self.args.train_encoder):
if WITH_LOSS_ZREG:
loss_zreg = torch.abs(z_sample).mean()
loss += loss_zreg
iter_stats.update({"loss_zreg": loss_zreg.item()})
encoding = self.update_encoding(z_sample)
iter_stats.update(encoding)
iter_stats["z_recon_mean"] = z_sample.mean().item()
iter_stats["z_recon_std"] = z_sample.std().item()
#######################
# Decoding
#######################
if not self.args.train_encoder:
z_sample = z_sample.detach()
with torch.set_grad_enabled(self.args.train_decoder):
# reconstruct images
X_recon = self.saae.P(z_sample)
#######################
# Reconstruction loss
#######################
loss_recon = aae_training.loss_recon(X_target, X_recon)
loss = loss_recon * self.args.w_rec
iter_stats["loss_recon"] = loss_recon.item()
#######################
# Structural loss
#######################
time_ssim_start = time.time()
cs_error_maps = None
if self.args.with_ssim_loss or eval:
store_cs_maps = (self._is_printout_iter(eval)
or eval) # get error maps for visualization
loss_ssim, cs_error_maps = aae_training.loss_struct(
X_target,
X_recon,
self.ssim,
calc_error_maps=store_cs_maps)
loss_ssim *= self.args.w_ssim
loss = 0.5 * loss + 0.5 * loss_ssim
iter_stats["ssim_torch"] = loss_ssim.item()
iter_stats.update({"time_ssim": time.time() - time_ssim_start})
#######################
# Adversarial loss
#######################
if (self.args.with_gan and self.args.train_decoder
and self.iter_in_epoch % 1 == 0):
gan_stats, loss_G = self.update_gan(
X_target,
X_recon,
z_sample,
train=not eval,
with_gen_loss=self.args.with_gen_loss,
)
loss += loss_G
iter_stats.update(gan_stats)
iter_stats["loss"] = loss.item()
if self.args.train_decoder:
loss.backward()
# Update auto-encoder
if not eval:
if self.args.train_encoder:
self.optimizer_E.step()
if self.args.train_decoder:
self.optimizer_G.step()
if eval or self._is_printout_iter(eval):
iter_stats["ssim"] = aae_training.calc_ssim(X_target, X_recon)
# statistics
iter_stats.update({
"epoch": self.epoch,
"timestamp": time.time(),
"iter_time": time.time() - self.iter_starttime,
"time_processing": time.time() - time_proc_start,
"iter": self.iter_in_epoch,
"total_iter": self.total_iter,
"batch_size": len(batch),
})
self.iter_starttime = time.time()
self.epoch_stats.append(iter_stats)
# print stats every N mini-batches
if self._is_printout_iter(eval):
self._print_iter_stats(
self.epoch_stats[-self._print_interval(eval):])
#
# Batch visualization
#
if self.args.show:
num_sample_images = {
128: 8,
256: 7,
512: 2,
1024: 1,
}
nimgs = num_sample_images[self.args.input_size]
self.visualize_random_images(nimgs, z_real=z_sample)
self.visualize_interpolations(z_sample, nimgs=2)
self.visualize_batch(
batch,
X_recon,
nimgs=nimgs,
ssim_maps=cs_error_maps,
ds=ds,
wait=self.wait,
)