Add lr drop

martinnovaak · martinnovaak · commit 142d17b69b44 · 2024-05-05T21:16:25.000+02:00
diff --git a/trainer/config.json b/trainer/config.json
@@ -6,6 +6,7 @@
   "epochs": 10,
   "batch_size": 16384,
   "wdl": 0.5,
+  "lr_drop_steps" : 4,
   "scale": 400,
   "hidden_layer_size": 16
 }
diff --git a/trainer/main.py b/trainer/main.py
@@ -25,6 +25,7 @@ def load_config(config_path="config.json"):
     epochs = config.get("epochs")
     batch_size = config.get("batch_size", 16384)
     wdl = config.get("wdl", 0.5)
+    lr_drop_steps = config.get("lr_drop_steps", 10)
     scale = config.get("scale")
 
     device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
@@ -38,7 +39,7 @@ def load_config(config_path="config.json"):
 
     start_time = time()
 
-    train(model, optimizer, dataloader, epochs, device)
+    train(model, optimizer, dataloader, epochs, lr_drop_steps, device)
 
     end_time = time()
     elapsed_time = end_time - start_time
diff --git a/trainer/train.py b/trainer/train.py
@@ -11,7 +11,7 @@ def print_epoch_stats(epoch, running_loss, iterations, fens, start_time, current
                .format(epoch, epoch_time, running_loss.item() / iterations, fens / epoch_time))
     print(message)
 
-def train(model: torch.nn.Module, optimizer: torch.optim.Optimizer, dataloader: BatchLoader, epochs: int, device: torch.device):
+def train(model: torch.nn.Module, optimizer: torch.optim.Optimizer, dataloader: BatchLoader, epochs: int, lr_drop_steps: int, device: torch.device):
     running_loss = torch.zeros(1, device=device)
     epoch_start_time = time()
     iterations = 0
@@ -23,6 +23,9 @@ def train(model: torch.nn.Module, optimizer: torch.optim.Optimizer, dataloader:
         if new_epoch:
             epoch += 1
 
+            if epoch % lr_drop_steps == 0:
+                optimizer.param_groups[0]["lr"] *= 0.1
+
             current_time = time()
             print_epoch_stats(epoch, running_loss, iterations, fens, epoch_start_time, current_time)
 

Original file line number	Diff line number	Diff line change
`@@ -6,6 +6,7 @@`
`6`	`6`	`"epochs": 10,`
`7`	`7`	`"batch_size": 16384,`
`8`	`8`	`"wdl": 0.5,`
	`9`	`+ "lr_drop_steps" : 4,`
`9`	`10`	`"scale": 400,`
`10`	`11`	`"hidden_layer_size": 16`
`11`	`12`	`}`