Merge branch 'master' into python3.8

# Conflicts:
#	models/hpm.py
#	models/layers.py
#	models/model.py
#	models/rgb_part_net.py
#	utils/configuration.py

1 files changed, 44 insertions, 14 deletions

diff --git a/models/model.py b/models/model.py
index 1dd195e..5af2d76 100644
--- a/models/model.py
+++ b/models/model.py
@@ -54,7 +54,6 @@ class Model:
         self.total_iters = self.meta.get('total_iters', (80000, 80000, 80000))
 
         self.is_train: bool = True
-        self.train_size: int = 74
         self.in_channels: int = 3
         self.pr: Optional[int] = None
         self.k: Optional[int] = None
@@ -70,6 +69,7 @@ class Model:
         self.optimizer: Optional[optim.Adam] = None
         self.scheduler: Optional[optim.lr_scheduler.StepLR] = None
         self.writer: Optional[SummaryWriter] = None
+        self.image_log_on = system_config.get('image_log_on', False)
 
         self.CASIAB_GALLERY_SELECTOR = {
             'selector': {'conditions': ClipConditions({r'nm-0[1-4]'})}
@@ -141,19 +141,18 @@ class Model:
         # Prepare for model, optimizer and scheduler
         model_hp = self.hp.get('model', {})
         optim_hp: Dict = self.hp.get('optimizer', {}).copy()
+        start_iter = optim_hp.pop('start_iter', 0)
         ae_optim_hp = optim_hp.pop('auto_encoder', {})
         pn_optim_hp = optim_hp.pop('part_net', {})
         hpm_optim_hp = optim_hp.pop('hpm', {})
         fc_optim_hp = optim_hp.pop('fc', {})
         sched_hp = self.hp.get('scheduler', {})
-        self.rgb_pn = RGBPartNet(self.train_size, self.in_channels, **model_hp)
+        self.rgb_pn = RGBPartNet(self.in_channels, **model_hp,
+                                 image_log_on=self.image_log_on)
         # Try to accelerate computation using CUDA or others
         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)
         self.writer = SummaryWriter(self._log_name)
@@ -171,10 +170,20 @@ class Model:
 
         # Training start
         start_time = datetime.now()
-        running_loss = torch.zeros(4).to(self.device)
+        running_loss = torch.zeros(5, device=self.device)
         print(f"{'Iter':^5} {'Loss':^6} {'Xrecon':^8} {'PoseSim':^8}",
-              f"{'CanoCons':^8} {'BATrip':^8} {'LR':^9}")
+              f"{'CanoCons':^8} {'BATripH':^8} {'BATripP':^8} LR(s)")
         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()
@@ -182,12 +191,10 @@ class Model:
             x_c1 = batch_c1['clip'].to(self.device)
             x_c2 = batch_c2['clip'].to(self.device)
             y = batch_c1['label'].to(self.device)
-            losses = self.rgb_pn(x_c1, x_c2, y)
+            losses, images = self.rgb_pn(x_c1, x_c2, y)
             loss = losses.sum()
             loss.backward()
             self.optimizer.step()
-            # Step scheduler
-            self.scheduler.step()
 
             # Statistics and checkpoint
             running_loss += losses.detach()
@@ -195,15 +202,39 @@ class Model:
             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'
+                '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()
                 print(f'{self.curr_iter:5d} {running_loss.sum() / 100:6.3f}',
-                      '{:f} {:f} {:f} {:f}'.format(*running_loss / 100),
-                      f'{self.scheduler.get_last_lr()[0]:.3e}')
+                      '{:f} {:f} {:f} {:f} {:f}'.format(*running_loss / 100),
+                      ' '.join(('{:.3e}'.format(lr) for lr in lrs)))
                 running_loss.zero_()
 
+            # Step scheduler
+            self.scheduler.step()
+
             if self.curr_iter % 1000 == 0:
                 torch.save({
                     'iter': self.curr_iter,
@@ -399,7 +430,6 @@ class Model:
             self,
             dataset_config: DatasetConfiguration
     ) -> Union[CASIAB]:
-        self.train_size = dataset_config.get('train_size', 74)
         self.in_channels = dataset_config.get('num_input_channels', 3)
         self._dataset_sig = self._make_signature(
             dataset_config,