Training stage seems to work, creating a non-run commit to use colab as an orchestration machine

dvc.yaml

@@ -8,3 +8,10 @@ stages:
     - src/data/raw/splits.mat
     outs:
     - src/data/processed/
+  train:
+    cmd: python3 src/code/training.py src/data/processed src/models
+    deps:
+    - src/code/training.py
+    - src/data/processed/
+    outs:
+    - src/models/

src/code/training.py

@@ -13,13 +13,13 @@ def get_y_fn(x):
 
 def create_data(data_path):
     fnames = get_files(data_path/'train', extensions='.jpg')
-    data = SegmentationDataLoaders.from_label_func(data_path/'train',  bs=1, num_workers=0, fnames=fnames, label_func=get_y_fn)
+    data = SegmentationDataLoaders.from_label_func(data_path/'train',  bs=4, num_workers=0, fnames=fnames, label_func=get_y_fn)
     return data
 
 
 def train(data):
-    learner = unet_learner(data, resnet34, metrics=rmse, wd=1e-2, n_out=1, loss_func=torch.nn.SmoothL1Loss())
-    learner.fine_tune(8)
+    learner = unet_learner(data, resnet34, metrics=rmse, wd=1e-2, n_out=1, loss_func=MSELossFlat())
+    learner.fine_tune(1)
 
 
 if __name__ == "__main__":