Don't run all epochs in test mode

PolicyEngine · Dec 22, 2024 · ef50800 · ef50800
1 parent d0c8c85
commit ef50800
Show file tree

Hide file tree

Showing 4 changed files with 10 additions and 5 deletions.
diff --git a/.github/workflows/pull_request.yaml b/.github/workflows/pull_request.yaml
@@ -49,6 +49,8 @@ jobs:
           HUGGING_FACE_TOKEN: ${{ secrets.HUGGING_FACE_TOKEN }}
       - name: Build datasets
         run: make data
+        env:
+          TEST_LITE: true
       - name: Run tests
         run: pytest
       - name: Test documentation builds

diff --git a/policyengine_uk_data/datasets/frs/local_areas/constituencies/calibrate.py b/policyengine_uk_data/datasets/frs/local_areas/constituencies/calibrate.py
@@ -4,6 +4,7 @@
 import numpy as np
 from tqdm import tqdm
 import h5py
+import os
 from policyengine_uk_data.datasets.frs.local_areas.constituencies.transform_constituencies import (
     transform_2010_to_2024,
 )
@@ -69,7 +70,7 @@ def dropout_weights(weights, p):
 
     optimizer = torch.optim.Adam([weights], lr=0.05)
 
-    desc = range(2048)
+    desc = range(128) if os.environ.get("DATA_LITE") else range(2048)
 
     for epoch in desc:
         optimizer.zero_grad()
@@ -78,7 +79,7 @@ def dropout_weights(weights, p):
         l.backward()
         optimizer.step()
         if epoch % 50 == 0:
-            print(f"Loss: {l.item()}, Epoch: {epoch}")
+            print(f"Loss: {l.item()}, Epoch: {epoch}", flush=True)
 
     final_weights = torch.exp(weights).detach().numpy()
     mapping_matrix = pd.read_csv(

diff --git a/policyengine_uk_data/datasets/frs/local_areas/local_authorities/calibrate.py b/policyengine_uk_data/datasets/frs/local_areas/local_authorities/calibrate.py
@@ -4,10 +4,11 @@
 import numpy as np
 from tqdm import tqdm
 import h5py
+import os
 from policyengine_uk_data.storage import STORAGE_FOLDER
 
 
-from loss import (
+from policyengine_uk_data.datasets.frs.local_areas.local_authorities.loss import (
     create_local_authority_target_matrix,
     create_national_target_matrix,
 )
@@ -62,7 +63,7 @@ def dropout_weights(weights, p):
 
     optimizer = torch.optim.Adam([weights], lr=0.05)
 
-    desc = range(2048)
+    desc = range(128) if os.environ.get("DATA_LITE") else range(2048)
 
     for epoch in desc:
         optimizer.zero_grad()

diff --git a/policyengine_uk_data/utils/reweight.py b/policyengine_uk_data/utils/reweight.py
@@ -1,5 +1,6 @@
 import numpy as np
 import torch
+import os
 
 
 def reweight(
@@ -47,7 +48,7 @@ def dropout_weights(weights, p):
 
     start_loss = None
 
-    iterator = range(1_000)
+    iterator = range(128) if os.environ.get("DATA_LITE") else range(2048)
     for i in iterator:
         optimizer.zero_grad()
         weights_ = dropout_weights(weights, dropout_rate)