Enable optimizer fine tuning

4 files changed, 47 insertions, 15 deletions

diff --git a/config.py b/config.py
index 791e888..c93fe19 100644
--- a/config.py
+++ b/config.py
@@ -33,7 +33,7 @@ config: Configuration = {
         # Batch size (pr, k)
         # `pr` denotes number of persons
         # `k` denotes number of sequences per person
-        'batch_size': (8, 16),
+        'batch_size': (2, 2),
         # Number of workers of Dataloader
         'num_workers': 4,
         # Faster data transfer from RAM to GPU if enabled
@@ -69,13 +69,23 @@ config: Configuration = {
             'triplet_margin': 0.2,
         },
         'optimizer': {
+            # Global parameters
             # Initial learning rate of Adam Optimizer
             'lr': 1e-4,
             # Coefficients used for computing running averages of
             #   gradient and its square
-            'betas': (0.9, 0.999),
+            # 'betas': (0.9, 0.999),
+            # Term added to the denominator
+            # 'eps': 1e-8,
             # Weight decay (L2 penalty)
-            'weight_decay': 0.001,
+            # 'weight_decay': 0,
+            # Use AMSGrad or not
+            # 'amsgrad': False,
+
+            # Local parameters (override global ones)
+            'auto_encoder': {
+                'weight_decay': 0.001
+            },
         },
         'scheduler': {
             # Period of learning rate decay
diff --git a/models/model.py b/models/model.py
index 740cdf3..8797636 100644
--- a/models/model.py
+++ b/models/model.py
@@ -133,12 +133,21 @@ class Model:
         dataloader = self._parse_dataloader_config(dataset, dataloader_config)
         # Prepare for model, optimizer and scheduler
         model_hp = self.hp.get('model', {})
-        optim_hp = self.hp.get('optimizer', {})
+        optim_hp: dict = self.hp.get('optimizer', {}).copy()
+        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)
         # Try to accelerate computation using CUDA or others
         self.rgb_pn = self.rgb_pn.to(self.device)
-        self.optimizer = optim.Adam(self.rgb_pn.parameters(), **optim_hp)
+        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)
 
diff --git a/test/model.py b/test/model.py
index f7fc57e..6fbcf00 100644
--- a/test/model.py
+++ b/test/model.py
@@ -14,12 +14,12 @@ def test_default_signature():
     model._parse_dataloader_config(casiab, conf['dataloader'])
     assert model._log_name == os.path.join(
         'runs', 'logs', 'RGB-GaitPart_80000_64_128_128_64_1_2_4_True_True_32_5_'
-                        '3_3_3_3_3_2_1_1_1_1_1_0_2_3_4_16_256_0.2_0.0001_0.9_'
-                        '0.999_0.001_500_0.9_CASIA-B_74_30_15_3_64_32_8_16')
+                        '3_3_3_3_3_2_1_1_1_1_1_0_2_3_4_16_256_0.2_0.0001_0.001_'
+                        '500_0.9_CASIA-B_74_30_15_3_64_32_8_16')
     assert model._checkpoint_sig == ('RGB-GaitPart_0_80000_64_128_128_64_1_2_4_'
                                      'True_True_32_5_3_3_3_3_3_2_1_1_1_1_1_0_2_'
-                                     '3_4_16_256_0.2_0.0001_0.9_0.999_0.001_'
-                                     '500_0.9_CASIA-B_74_30_15_3_64_32_8_16')
+                                     '3_4_16_256_0.2_0.0001_0.001_500_0.9_'
+                                     'CASIA-B_74_30_15_3_64_32_8_16')
 
 
 def test_default_signature_with_selector():
@@ -31,11 +31,10 @@ def test_default_signature_with_selector():
     model._parse_dataloader_config(casiab, conf['dataloader'])
     assert model._log_name == os.path.join(
         'runs', 'logs', 'RGB-GaitPart_80000_64_128_128_64_1_2_4_True_True_32_5_'
-                        '3_3_3_3_3_2_1_1_1_1_1_0_2_3_4_16_256_0.2_0.0001_0.9_'
-                        '0.999_0.001_500_0.9_CASIA-B_74_30_15_3_64_32_bg-0\\d_'
-                        'nm-0\\d_8_16')
+                        '3_3_3_3_3_2_1_1_1_1_1_0_2_3_4_16_256_0.2_0.0001_0.001_'
+                        '500_0.9_CASIA-B_74_30_15_3_64_32_bg-0\\d_nm-0\\d_8_16')
     assert model._checkpoint_sig == ('RGB-GaitPart_0_80000_64_128_128_64_1_2_4_'
                                      'True_True_32_5_3_3_3_3_3_2_1_1_1_1_1_0_2_'
-                                     '3_4_16_256_0.2_0.0001_0.9_0.999_0.001_'
-                                     '500_0.9_CASIA-B_74_30_15_3_64_32_bg-0\\d_'
-                                     'nm-0\\d_8_16')
+                                     '3_4_16_256_0.2_0.0001_0.001_500_0.9_'
+                                     'CASIA-B_74_30_15_3_64_32_bg-0\\d_nm-0\\d_'
+                                     '8_16')
diff --git a/utils/configuration.py b/utils/configuration.py
index 9a8c2ae..8b265e8 100644
--- a/utils/configuration.py
+++ b/utils/configuration.py
@@ -43,10 +43,24 @@ class ModelHPConfiguration(TypedDict):
     triplet_margin: float
 
 
+class SubOptimizerHPConfiguration(TypedDict):
+    lr: int
+    betas: tuple[float, float]
+    eps: float
+    weight_decay: float
+    amsgrad: bool
+
+
 class OptimizerHPConfiguration(TypedDict):
     lr: int
     betas: tuple[float, float]
+    eps: float
     weight_decay: float
+    amsgrad: bool
+    auto_encoder: SubOptimizerHPConfiguration
+    part_net: SubOptimizerHPConfiguration
+    hpm: SubOptimizerHPConfiguration
+    fc: SubOptimizerHPConfiguration
 
 
 class SchedulerHPConfiguration(TypedDict):