Merge remote-tracking branch 'origin/main' into method_args

.github/workflows/auto-publish.yml

@@ -12,37 +12,32 @@ jobs:
         - name: Set up Python
           uses: actions/[email protected]
           with:
-            python-version: 3.8
+            python-version: 3.9
 
         - name: Install poetry
           run: make setup
-
+
+        - name: Install poeblix for version freezing
+          run: poetry self add poeblix@latest
+
         - name: Bump version
           run: |
             poetry version prerelease
-    
 
         - name: Commit changes
           run: |
             git config --local user.email "[email protected]"
             git config --local user.name "GitHub Action"
             git add .
             git commit -m "Bump version" || echo "No changes to commit"
-   
-
-        - name: Update diagram
-          uses: githubocto/repo-visualizer@main
-          with:
-            excluded_paths: "ignore,.github"
-            should_push: false
-
+
         - name: Push changes
           run: git push
+
+        - name: Build with lockfile versions
+          run: poetry blixbuild --only-lock
 
-        - name: Build and publish to pypi
-          uses: JRubics/poetry-publish@v1.16
+        - name: Publish distribution 📦 to PyPI
+          uses: pypa/gh-action-pypi-publish@release/v1
           with:
-            python_version: 3.8
-            pypi_token: ${{ secrets.PYPI_TOKEN }}
-            allow_poetry_pre_release: "yes"
-            ignore_dev_requirements: "yes"
+            password: ${{ secrets.PYPI_TOKEN }}

.github/workflows/build.yml

@@ -1,13 +1,19 @@
 name: Lighter CI
 
-on: [push, pull_request]
-
+on:
+  push:
+    branches:
+      - main
+  pull_request:
+    branches:
+      - main
+
 jobs:
   build:
     runs-on: ubuntu-latest
     strategy:
       matrix:
-        python-version: ["3.8", "3.9", "3.10"]
+        python-version: ["3.9", "3.10", "3.11"]
 
     steps:
     - uses: actions/checkout@v2

docs/basics/quickstart.md

@@ -70,6 +70,13 @@ system:
           - _target_: torchvision.transforms.Normalize
             mean: [0.5, 0.5, 0.5]
             std: [0.5, 0.5, 0.5]
+
+  postprocessing:
+      # Format the batch as required by Lighter as lighter 
+      # takes a dictionary as input. 
+      batch:
+          train: '$lambda x: {"input": x[0], "target": x[1]}'
+
 ```
 
 For more information about each of the LighterSystem components and how to override them, see [here](./config.md)
@@ -79,6 +86,10 @@ We just combine the Trainer and LighterSystem into a single YAML and run the com
 
 === "cifar10.yaml"
     ```yaml
+    trainer:
+      _target_: pytorch_lightning.Trainer
+      max_epochs: 100
+
     system:
       _target_: lighter.LighterSystem
       batch_size: 512
@@ -109,6 +120,12 @@ We just combine the Trainer and LighterSystem into a single YAML and run the com
                 mean: [0.5, 0.5, 0.5]
                 std: [0.5, 0.5, 0.5]
 
+      postprocessing:
+          # Format the batch as required by Lighter as lighter 
+          # takes a dictionary as input. 
+          batch:
+              train: '$lambda x: {"input": x[0], "target": x[1]}'
+
     ```
 === "Terminal"
     ```

lighter/__init__.py

@@ -1,4 +1,4 @@
-__version__ = "0.0.2a25"
+__version__ = "0.0.2a29"
 
 from .utils.logging import _setup_logging
 

lighter/system.py

@@ -205,9 +205,12 @@ def _base_step(self, batch: Dict, batch_idx: int, mode: str) -> Union[Dict[str,
         target = apply_fns(target, self.postprocessing["logging"]["target"])
         pred = apply_fns(pred, self.postprocessing["logging"]["pred"])
 
-        # If the loss is a dict, the sublosses must be combined under "total" key.
+        # Ensure that a dict of losses has a 'total' key.
         if isinstance(loss, dict) and "total" not in loss:
-            raise ValueError("The loss dictionary must have 'total' loss, combining all the sublosses.")
+            raise ValueError(
+                "The loss dictionary must include a 'total' key that combines all sublosses. "
+                "Example: {'total': combined_loss, 'subloss1': loss1, ...}"
+            )
 
         # Logging
         self._log_stats(loss, metrics, mode, batch_idx)
@@ -251,6 +254,8 @@ def _log_stats(
         # Metrics
         if metrics is not None:
             for name, metric in metrics.items():
+                if not isinstance(metric, Metric):
+                    raise TypeError(f"Expected type for metric is 'Metric', got '{type(metric).__name__}' instead.")
                 on_step_log(f"{mode}/metrics/{name}/step", metric)
                 on_epoch_log(f"{mode}/metrics/{name}/epoch", metric)
         # Optimizer's lr, momentum, beta. Logged in train mode and once per epoch.