First finalized pipeline

TODO: Add checkpointing capability

dvc.lock

@@ -3,8 +3,8 @@ process_data:
     src/data/processed
   deps:
   - path: src/code/make_dataset.py
-    md5: 3a0dee3a1ba9c587b8ca6ea6f0447ada
-    size: 5227
+    md5: e069c7323c9be16baedd8f988375e145
+    size: 5256
   - path: src/data/raw/nyu_depth_v2_labeled.mat
     md5: 520609c519fba3ba5ac58c8fefcc3530
     size: 2972037809

dvc.yaml

@@ -11,21 +11,26 @@ stages:
   train:
     cmd: python3 src/code/training.py src/data/processed/train
     deps:
+    - src/code/custom_data_loading.py
+    - src/code/params.yml
     - src/code/training.py
     - src/data/processed/train
     outs:
     - src/models/
+    - logs/train_params.yml:
+        cache: false
     metrics:
     - logs/train_metrics.csv:
         cache: false
   eval:
     cmd: python3 src/code/eval.py src/data/processed/test
     deps:
+    - src/code/params.yml
+    - src/code/custom_data_loading.py
+    - src/code/eval_metric_calculation.py
     - src/code/eval.py
     - src/models/model.pth
     - src/data/processed/test
-    outs:
-    - src/eval/
     metrics:
     - logs/test_metrics.csv:
         cache: false

src/code/custom_data_loading.py

@@ -1,3 +1,4 @@
+import yaml
 from fastai.vision.all import \
     DataLoaders, \
     delegates, \
@@ -7,7 +8,8 @@ from fastai.vision.all import \
     PILImageBW, \
     RandomSplitter, \
     Path, \
-    get_files
+    get_files, \
+    L
 
 
 class ImageImageDataLoaders(DataLoaders):
@@ -33,13 +35,24 @@ def get_y_fn(x):
     return y
 
 
-def create_data(data_path):
+def create_data(data_path, is_test=False):
+    with open(r"./src/code/params.yml") as f:
+        params = yaml.safe_load(f)
+
     filenames = get_files(data_path, extensions='.jpg')
     if len(filenames) == 0:
         raise ValueError("Could not find any files in the given path")
     dataset = ImageImageDataLoaders.from_label_func(data_path,
-                                                    seed=42,
-                                                    bs=4, num_workers=0,
+                                                    seed=int(params['seed']),
+                                                    bs=int(params['batch_size']),
+                                                    num_workers=int(params['num_workers']),
                                                     filenames=filenames,
                                                     label_func=get_y_fn)
+
+    if is_test:
+        filenames = get_files(Path(data_path), extensions='.jpg')
+        test_files = L([Path(i) for i in filenames])
+        test_dl = dataset.test_dl(test_files, with_labels=True)
+        return dataset, test_dl
+
     return dataset

src/code/eval.py

@@ -1,7 +1,8 @@
 import sys
-from fastai.vision.all import unet_learner, Path, resnet34, MSELossFlat, get_files, L, tuplify
-from src.code.custom_data_loading import create_data
-from src.code.eval_metric_calculation import compute_eval_metrics
+import yaml
+from fastai.vision.all import unet_learner, Path, resnet34, MSELossFlat, tuplify
+from custom_data_loading import create_data
+from eval_metric_calculation import compute_eval_metrics
 from dagshub import dagshub_logger
 
 
@@ -10,19 +11,24 @@ if __name__ == "__main__":
         print("usage: %s <test_data_path>" % sys.argv[0], file=sys.stderr)
         sys.exit(0)
 
+    with open(r"./src/code/params.yml") as f:
+        params = yaml.safe_load(f)
+
     data_path = Path(sys.argv[1])
-    data = create_data(data_path)
+    data, test_dl = create_data(data_path, is_test=True)
+
+    arch = {'resnet34': resnet34}
+    loss = {'MSELossFlat': MSELossFlat()}
 
-    filenames = get_files(Path(sys.argv[1]), extensions='.jpg')
-    test_files = L([Path(i) for i in filenames])
-    test_dl = data.test_dl(test_files, with_labels=True)
     learner = unet_learner(data,
-                           resnet34,
-                           n_out=3,
-                           loss_func=MSELossFlat(),
+                           arch.get(params['architecture']),
+                           n_out=int(params['num_outs']),
+                           loss_func=loss.get(params['loss_func']),
                            path='src/',
                            model_dir='models')
     learner = learner.load('model')
+
+    print("Running model on test data...")
     inputs, predictions, targets, decoded = learner.get_preds(dl=test_dl,
                                                               with_input=True,
                                                               with_decoded=True)
@@ -31,6 +37,7 @@ if __name__ == "__main__":
     decoded_predictions = learner.dls.decode(inputs + tuplify(decoded))[1]
     decoded_targets = learner.dls.decode(inputs + tuplify(targets))[1]
 
+    print("Calculating metrics...")
     metrics = compute_eval_metrics(decoded_targets.numpy(), decoded_predictions.numpy())
 
     with dagshub_logger(

src/code/params.yml

@@ -0,0 +1,13 @@
+seed: 42
+data: nyu_depth_v2
+batch_size: 4
+num_workers: 0
+weight_decay: 1e-2
+learning_rate: 1e-3
+epochs: 1
+num_outs: 3
+source_dir: src
+model_dir: models
+architecture: resnet34
+loss_func: MSELossFlat
+train_metric: rmse

src/code/training.py

@@ -3,8 +3,9 @@ Receives 1 arguments from argparse:
   <data_path> - Path to the dataset which is split into 2 folders - train and test.
 """
 import sys
+import yaml
 from fastai.vision.all import unet_learner, Path, resnet34, rmse, MSELossFlat
-from src.code.custom_data_loading import create_data
+from custom_data_loading import create_data
 from dagshub.fastai import DAGsHubLogger
 
 
@@ -14,23 +15,30 @@ if __name__ == "__main__":
         print("usage: %s <data_path>" % sys.argv[0], file=sys.stderr)
         sys.exit(0)
 
+    with open(r"./src/code/params.yml") as f:
+        params = yaml.safe_load(f)
+
     data = create_data(Path(sys.argv[1]))
-    wd, lr, ep = 1e-2, 1e-3, 1
+
+    metrics = {'rmse': rmse}
+    arch = {'resnet34': resnet34}
+    loss = {'MSELossFlat': MSELossFlat()}
+
     learner = unet_learner(data,
-                           resnet34,
-                           metrics=rmse,
-                           wd=wd,
-                           n_out=3,
-                           loss_func=MSELossFlat(),
-                           path='src/',
-                           model_dir='models',
+                           arch.get(params['architecture']),
+                           metrics=metrics.get(params['train_metric']),
+                           wd=float(params['weight_decay']),
+                           n_out=int(params['num_outs']),
+                           loss_func=loss.get(params['loss_func']),
+                           path=params['source_dir'],
+                           model_dir=params['model_dir'],
                            cbs=DAGsHubLogger(
                                metrics_path="logs/train_metrics.csv",
-                               hparams_path="logs/train_params.yml"
-                           ))
+                               hparams_path="logs/train_params.yml"))
 
     print("Training model...")
-    learner.fine_tune(epochs=ep, base_lr=lr)
+    learner.fine_tune(epochs=int(params['epochs']),
+                      base_lr=float(params['learning_rate']))
     print("Saving model...")
     learner.save('model')
     print("Done!")