add a warning

mwalmsley · Mar 2, 2024 · eaa98ce · eaa98ce
1 parent ce2c80a
commit eaa98ce
Showing 1 changed file with 5 additions and 3 deletions.
diff --git a/zoobot/pytorch/training/finetune.py b/zoobot/pytorch/training/finetune.py
@@ -112,14 +112,13 @@ def __init__(
             # self.encoder_dim = 9216
         else:
             self.encoder_dim = define_model.get_encoder_dim(self.encoder)
-        self.n_blocks = n_blocks
-        logging.info('Blocks to finetune: {}'.format(n_layers))
 
         # for backwards compat.
         if n_layers:
             logging.warning('FinetuneableZoobot(n_layers) is now renamed to n_blocks, please update to pass n_blocks instead! For now, setting n_blocks=n_layers')
             self.n_blocks = n_layers
-            logging.info('Layers to finetune: {}'.format(n_layers))
+        else:
+            self.n_blocks = n_blocks
 
         self.learning_rate = learning_rate
         self.lr_decay = lr_decay
@@ -243,13 +242,16 @@ def configure_optimizers(self):
 
 
         logging.info('param groups: {}'.format(len(params)))
+
+        # because it iterates through the generators, THIS BREAKS TRAINING so only uncomment to debug params
         # for param_group_n, param_group in enumerate(params):
         #     shapes_within_param_group = [p.shape for p in list(param_group['params'])]
         #     logging.debug('param group {}: {}'.format(param_group_n, shapes_within_param_group))
         # print('head params to optimize', [p.shape for p in params[0]['params']])  # head only
         # print(list(param_group['params']) for param_group in params)
         # exit()
         # Initialize AdamW optimizer
+
         opt = torch.optim.AdamW(params, weight_decay=self.weight_decay)  # lr included in params dict
         logging.info('Optimizer ready, configuring scheduler')