From 9f3d4cc14ad36e515b56e86fb8e26f519bde831e Mon Sep 17 00:00:00 2001 From: Jordan Gong Date: Tue, 9 Feb 2021 16:51:08 +0800 Subject: Some optimizations 1. Scheduler will decay the learning rate of auto-encoder only 2. Write learning rate history to tensorboard 3. Reduce image log frequency --- models/model.py | 76 +++++++++++++++++++++++++++++++++------------------------ 1 file changed, 44 insertions(+), 32 deletions(-) (limited to 'models/model.py') diff --git a/models/model.py b/models/model.py index 0418070..ee07615 100644 --- a/models/model.py +++ b/models/model.py @@ -153,8 +153,18 @@ class Model: self.rgb_pn = self.rgb_pn.to(self.device) self.optimizer = optim.Adam([ {'params': self.rgb_pn.ae.parameters(), **ae_optim_hp}, + {'params': self.rgb_pn.pn.parameters(), **pn_optim_hp}, + {'params': self.rgb_pn.hpm.parameters(), **hpm_optim_hp}, + {'params': self.rgb_pn.fc_mat, **fc_optim_hp} ], **optim_hp) - self.scheduler = optim.lr_scheduler.StepLR(self.optimizer, **sched_hp) + sched_gamma = sched_hp.get('gamma', 0.9) + sched_step_size = sched_hp.get('step_size', 500) + self.scheduler = optim.lr_scheduler.LambdaLR(self.optimizer, lr_lambda=[ + lambda epoch: sched_gamma ** (epoch // sched_step_size), + lambda epoch: 0 if epoch < start_iter else 1, + lambda epoch: 0 if epoch < start_iter else 1, + lambda epoch: 0 if epoch < start_iter else 1, + ]) self.writer = SummaryWriter(self._log_name) self.rgb_pn.train() @@ -172,18 +182,8 @@ class Model: start_time = datetime.now() running_loss = torch.zeros(5, device=self.device) print(f"{'Iter':^5} {'Loss':^6} {'Xrecon':^8} {'PoseSim':^8}", - f"{'CanoCons':^8} {'BATripH':^8} {'BATripP':^8} LR(s)") + f"{'CanoCons':^8} {'BATripH':^8} {'BATripP':^8} {'LRs':^19}") for (batch_c1, batch_c2) in dataloader: - if self.curr_iter == start_iter: - self.optimizer.add_param_group( - {'params': self.rgb_pn.pn.parameters(), **pn_optim_hp} - ) - self.optimizer.add_param_group( - {'params': self.rgb_pn.hpm.parameters(), **hpm_optim_hp} - ) - self.optimizer.add_param_group( - {'params': self.rgb_pn.fc_mat, **fc_optim_hp} - ) self.curr_iter += 1 # Zero the parameter gradients self.optimizer.zero_grad() @@ -205,31 +205,39 @@ class Model: 'Canonical consistency loss', 'Batch All triplet loss (HPM)', 'Batch All triplet loss (PartNet)' ], losses)), self.curr_iter) - if self.image_log_on: - (appearance_image, canonical_image, pose_image) = images - self.writer.add_images( - 'Canonical image', canonical_image, self.curr_iter - ) - for i in range(self.pr * self.k): - self.writer.add_images( - f'Original image/batch {i}', x_c1[i], self.curr_iter - ) - self.writer.add_images( - f'Appearance image/batch {i}', - appearance_image[:, i, :, :, :], - self.curr_iter - ) - self.writer.add_images( - f'Pose image/batch {i}', - pose_image[:, i, :, :, :], - self.curr_iter - ) if self.curr_iter % 100 == 0: lrs = self.scheduler.get_last_lr() + # Write learning rates + self.writer.add_scalar( + 'Learning rate/Auto-encoder', lrs[0], self.curr_iter + ) + self.writer.add_scalar( + 'Learning rate/Others', lrs[1], self.curr_iter + ) + # Write disentangled images + if self.image_log_on: + (appearance_image, canonical_image, pose_image) = images + self.writer.add_images( + 'Canonical image', canonical_image, self.curr_iter + ) + for i in range(self.pr * self.k): + self.writer.add_images( + f'Original image/batch {i}', x_c1[i], self.curr_iter + ) + self.writer.add_images( + f'Appearance image/batch {i}', + appearance_image[:, i, :, :, :], + self.curr_iter + ) + self.writer.add_images( + f'Pose image/batch {i}', + pose_image[:, i, :, :, :], + self.curr_iter + ) print(f'{self.curr_iter:5d} {running_loss.sum() / 100:6.3f}', '{:f} {:f} {:f} {:f} {:f}'.format(*running_loss / 100), - ' '.join(('{:.3e}'.format(lr) for lr in lrs))) + '{:.3e} {:.3e}'.format(lrs[0], lrs[1])) running_loss.zero_() # Step scheduler @@ -244,6 +252,10 @@ class Model: }, self._checkpoint_name) print(datetime.now() - start_time, 'used') start_time = datetime.now() + print( + f"{'Iter':^5} {'Loss':^6} {'Xrecon':^8} {'PoseSim':^8}", + f"{'CanoCons':^8} {'BATripH':^8} {'BATripP':^8} {'LRs':^19}" + ) if self.curr_iter == self.total_iter: self.writer.close() -- cgit v1.2.3 From 916cf90d04e57fee23092c966740fbe94fd92cff Mon Sep 17 00:00:00 2001 From: Jordan Gong Date: Tue, 9 Feb 2021 21:21:57 +0800 Subject: Improve performance when disentangling This is a HUGE performance optimization, up to 2x faster than before. Mainly because of the replacement of randomized for-loop with randomized tensor. --- models/model.py | 37 ++++++++++++++++--------------------- 1 file changed, 16 insertions(+), 21 deletions(-) (limited to 'models/model.py') diff --git a/models/model.py b/models/model.py index ee07615..70c43b3 100644 --- a/models/model.py +++ b/models/model.py @@ -181,8 +181,9 @@ class Model: # Training start start_time = datetime.now() running_loss = torch.zeros(5, device=self.device) - print(f"{'Iter':^5} {'Loss':^6} {'Xrecon':^8} {'PoseSim':^8}", - f"{'CanoCons':^8} {'BATripH':^8} {'BATripP':^8} {'LRs':^19}") + print(f"{'Time':^8} {'Iter':^5} {'Loss':^6}", + f"{'Xrecon':^8} {'CanoCons':^8} {'PoseSim':^8}", + f"{'BATripH':^8} {'BATripP':^8} {'LRs':^19}") for (batch_c1, batch_c2) in dataloader: self.curr_iter += 1 # Zero the parameter gradients @@ -201,8 +202,8 @@ class Model: # Write losses to TensorBoard self.writer.add_scalar('Loss/all', loss, self.curr_iter) self.writer.add_scalars('Loss/details', dict(zip([ - 'Cross reconstruction loss', 'Pose similarity loss', - 'Canonical consistency loss', 'Batch All triplet loss (HPM)', + 'Cross reconstruction loss', 'Canonical consistency loss', + 'Pose similarity loss', 'Batch All triplet loss (HPM)', 'Batch All triplet loss (PartNet)' ], losses)), self.curr_iter) @@ -217,25 +218,25 @@ class Model: ) # Write disentangled images if self.image_log_on: - (appearance_image, canonical_image, pose_image) = images + i_a, i_c, i_p = images self.writer.add_images( - 'Canonical image', canonical_image, self.curr_iter + 'Canonical image', i_c, self.curr_iter ) - for i in range(self.pr * self.k): + for (i, (o, a, p)) in enumerate(zip(x_c1, i_a, i_p)): self.writer.add_images( - f'Original image/batch {i}', x_c1[i], self.curr_iter + f'Original image/batch {i}', o, self.curr_iter ) self.writer.add_images( - f'Appearance image/batch {i}', - appearance_image[:, i, :, :, :], - self.curr_iter + f'Appearance image/batch {i}', a, self.curr_iter ) self.writer.add_images( - f'Pose image/batch {i}', - pose_image[:, i, :, :, :], - self.curr_iter + f'Pose image/batch {i}', p, self.curr_iter ) - print(f'{self.curr_iter:5d} {running_loss.sum() / 100:6.3f}', + time_used = datetime.now() - start_time + remaining_minute, second = divmod(time_used.seconds, 60) + hour, minute = divmod(remaining_minute, 60) + print(f'{hour:02}:{minute:02}:{second:02}', + f'{self.curr_iter:5d} {running_loss.sum() / 100:6.3f}', '{:f} {:f} {:f} {:f} {:f}'.format(*running_loss / 100), '{:.3e} {:.3e}'.format(lrs[0], lrs[1])) running_loss.zero_() @@ -250,12 +251,6 @@ class Model: 'optim_state_dict': self.optimizer.state_dict(), 'loss': loss, }, self._checkpoint_name) - print(datetime.now() - start_time, 'used') - start_time = datetime.now() - print( - f"{'Iter':^5} {'Loss':^6} {'Xrecon':^8} {'PoseSim':^8}", - f"{'CanoCons':^8} {'BATripH':^8} {'BATripP':^8} {'LRs':^19}" - ) if self.curr_iter == self.total_iter: self.writer.close() -- cgit v1.2.3