Implement simple training test.

.github/workflows/ci.yml

@@ -27,6 +27,7 @@ jobs:
         source .env/bin/activate
         make install
         make install-dev
+        make install-training
     - name: Unit tests
       run: |
         source .env/bin/activate

Makefile

@@ -2,6 +2,9 @@ install: ## [Local development] Upgrade pip, install requirements, install packa
 	python -m pip install -U pip
 	python -m pip install -e .
 
+install-training:
+	python -m pip install -r requirements-training.txt
+
 install-dev: ## [Local development] Install test requirements
 	python -m pip install -r requirements-test.txt
 

README.md

@@ -118,6 +118,9 @@ Install pip PyTorch as per https://pytorch.org/get-started/locally/
 
 Test can be run with `make install-dev` then `make test`
 
+`python -m pytest -x -s -v tests -k "training"` to run a specific test
+
+
 #### Other dependencies
 
 Install open_clip pacakge and remaining dependencies:

src/training/main.py

@@ -1,5 +1,6 @@
 import logging
 import os
+import sys
 import random
 from datetime import datetime
 
@@ -38,8 +39,8 @@ def random_seed(seed=42, rank=0):
     random.seed(seed + rank)
 
 
-def main():
-    args = parse_args()
+def main(args):
+    args = parse_args(args)
 
     if torch.cuda.is_available():
         # This enables tf32 on Ampere GPUs which is only 8% slower than
@@ -314,4 +315,4 @@ def copy_codebase(args):
 
 
 if __name__ == "__main__":
-    main()
+    main(sys.argv[1:])

src/training/params.py

@@ -10,7 +10,7 @@ def get_default_params(model_name):
         return {"lr": 5.0e-4, "beta1": 0.9, "beta2": 0.999, "eps": 1.0e-8}
 
 
-def parse_args():
+def parse_args(args):
     parser = argparse.ArgumentParser()
     parser.add_argument(
         "--train-data",
@@ -293,7 +293,7 @@ def parse_args():
     parser.add_argument(
         "--grad-clip-norm", type=float, default=None, help="Gradient clip."
     )
-    args = parser.parse_args()
+    args = parser.parse_args(args)
 
     # If some params are not passed, we use the default values based on model name.
     default_params = get_default_params(args.model)

tests/test_training_simple.py

@@ -0,0 +1,22 @@
+
+import torch
+from PIL import Image
+from training.main import main
+import pytest
+import os
+os.environ["CUDA_VISIBLE_DEVICES"] = ""
+
+def test_training():
+    main([
+    '--save-frequency', '1',
+    '--zeroshot-frequency', '1',
+    '--dataset-type', "synthetic",
+    '--train-num-samples', '16',
+    '--warmup', '1',
+    '--batch-size', '4',
+    '--lr', '1e-3',
+    '--wd', '0.1',
+    '--epochs', '1',
+    '--workers', '2',
+    '--model', 'RN50'
+    ])