freeze decoder

Author: franciscoeacosta

neurometry/curvature/grid-cells-curvature/models/xu_rnn/default_config.py

@@ -22,7 +22,7 @@
 ###-----TRAINING PARAMETERS-----###
 load_pretrain=True
 pretrain_path=os.path.join(os.getcwd(),"logs/rnn_isometry/20240418-180712/ckpt/model/checkpoint-step25000.pth")
-num_steps_train=100#7500  # 10000
+num_steps_train=10000#7500  # 10000
 lr_decay_from=10000
 steps_per_logging=20
 steps_per_large_logging=500  # 500
@@ -81,4 +81,4 @@
 
 ###-----RAY TUNE PARAMETERS-----###
 sweep_metric= "error_reencode"
-num_samples = 1#1000
+num_samples = 1000#1000

neurometry/curvature/grid-cells-curvature/models/xu_rnn/experiment.py

@@ -34,7 +34,7 @@ def __init__(self, rng, config: ml_collections.ConfigDict, device):
         self.model_config = model.GridCellConfig(**config.model)
         self.model = model.GridCell(self.model_config).to(device)
 
-        if config.train.freeze_decoder:
+        if config.model.freeze_decoder:
             logging.info("==== freeze decoder ====")
             for param in self.model.decoder.parameters():
                 param.requires_grad = False
@@ -69,8 +69,8 @@ def __init__(self, rng, config: ml_collections.ConfigDict, device):
             logging.info(f"Loading pretrain model from {ckpt_model_path}")
             ckpt = torch.load(ckpt_model_path, map_location=device)
             self.model.load_state_dict(ckpt["state_dict"])
-            logging.info("==== load pretrained optimizer ====")
-            self.optimizer.load_state_dict(ckpt["optimizer"])
+            # logging.info("==== load pretrained optimizer ====")
+            # self.optimizer.load_state_dict(ckpt["optimizer"])
             self.starting_step = ckpt["step"]
         else:
             self.starting_step = 1

neurometry/curvature/grid-cells-curvature/models/xu_rnn/model.py

@@ -9,6 +9,7 @@
 
 @dataclass
 class GridCellConfig:
+    freeze_decoder: bool
     trans_type: str
     num_grid: int
     num_neurons: int