Migrated to fastai2, creating the DataLoader now works but I'm stuck on not being able to change the batch_size or num_workers as the interface seems to have changed

.gitignore

@@ -5,3 +5,4 @@
 .workspace/
 aws/
 google-cloud-sdk
+__pycache__/

Makefile

@@ -27,7 +27,11 @@ else
 	@echo ">>> New virtual env created. Activate with:\nsource env/bin/activate ."
 endif
 
-requirements:
+load_requirements:
 	@echo ">>> Installing requirements. Make sure your virtual environment is activated."
 	$(PYTHON_INTERPRETER) -m pip install -U pip setuptools wheel
-	$(PYTHON_INTERPRETER) -m pip install -r requirements.txt
+	$(PYTHON_INTERPRETER) -m pip install -r requirements.txt
+
+save_requirements:
+	@echo ">>> Saving requirements."
+	pip list --format=freeze > requirements.txt

README.md

@@ -28,8 +28,8 @@ If you'd like to take part, please follow the guide.
 
         ```bash
         $ make env
-        $ conda activate savta_depth
-        $ make requirements
+        $ source activate savta_depth
+        $ make load_requirements
         ```
 
         **Note**: If you don't have a GPU you will need to install pytorch separately and then run make requirements. You can install pytorch for computers without a gpu with the following command:
@@ -49,7 +49,7 @@ If you'd like to take part, please follow the guide.
     * Freeze your virtualenv by typing in the terminal:
 
         ```bash
-        pip freeze > requirements.txt
+        $ make save_requirements
         ```
 
     * Push your code to DAGsHub, and your dvc managed files to your dvc remote. In order to setup a dvc remote please refer to [this guide](https://dagshub.com/docs/getting-started/set-up-remote-storage-for-data-and-models/).

requirements.txt

@@ -1,5 +1,6 @@
 appdirs==1.4.4
 atpublic==2.0
+backcall==0.2.0
 blis==0.4.1
 cachetools==4.1.1
 catalogue==1.0.0
@@ -37,6 +38,9 @@ grandalf==0.6
 h5py==2.10.0
 idna==2.10
 importlib-metadata==1.7.0
+ipython==7.17.0
+ipython-genutils==0.2.0
+jedi==0.17.2
 joblib==0.16.0
 jsonpath-ng==1.5.1
 kiwisolver==1.2.0
@@ -52,12 +56,18 @@ olefile==0.46
 opencv-python==4.4.0.42
 packaging==20.4
 pandas==1.1.1
+parso==0.7.1
 pathspec==0.8.0
+pexpect==4.8.0
+pickleshare==0.7.5
 Pillow==7.2.0
+pip==20.2.2
 plac==1.1.3
 ply==3.11
 preshed==3.0.2
+prompt-toolkit==3.0.6
 protobuf==3.13.0
+ptyprocess==0.6.0
 pyasn1==0.4.8
 pyasn1-modules==0.2.8
 pycparser==2.20
@@ -75,6 +85,7 @@ ruamel.yaml==0.16.10
 ruamel.yaml.clib==0.2.0
 scikit-learn==0.23.2
 scipy==1.5.2
+setuptools==49.6.0.post20200814
 shortuuid==1.0.1
 shtab==1.3.1
 six==1.15.0
@@ -88,9 +99,12 @@ toml==0.10.1
 torch==1.6.0
 torchvision==0.7.0
 tqdm==4.48.2
+traitlets==4.3.3
 typing-extensions==3.7.4.3
 urllib3==1.25.10
 voluptuous==0.11.7
 wasabi==0.7.1
+wcwidth==0.2.5
+wheel==0.35.1
 zc.lockfile==2.0
 zipp==3.1.0

run_dev_env.sh

@@ -1,7 +1,7 @@
 docker run -d \
     -p 8080:8080 \
-    --name "ml-workspace" -v "${PWD}:/workspace" \
+    --name "dags-ml-workspace" -v "${PWD}:/workspace" \
     --env AUTHENTICATE_VIA_JUPYTER="dagshub_savta" \
     --shm-size 2G \
     --restart always \
-    mltooling/ml-workspace-minimal:latest
+    dagshub/ml-workspace-minimal:latest

src/code/training.py

@@ -1,6 +1,7 @@
 import torch
 import sys
-from fastai.vision import unet_learner, ImageImageList, models, Path, root_mean_squared_error
+from fastai2.vision.all import *
+from torchvision.utils import save_image
 
 
 def get_y_fn(x):
@@ -10,16 +11,16 @@ def get_y_fn(x):
     return y
 
 
-def create_databunch(data_path):
-    data = (ImageImageList.from_folder(data_path)
-            .filter_by_func(lambda fname: fname.suffix == '.jpg')
-            .split_by_folder(train='train', valid='test')
-            .label_from_func(get_y_fn).databunch()).normalize()
+def create_data(data_path):
+    fnames = get_files(data_path/'train', extensions='.jpg')
+    data = SegmentationDataLoaders.from_label_func(data_path/'train', fnames=fnames, label_func=get_y_fn)
     return data
 
 
 def train(data):
-    learner = unet_learner(data, models.resnet34, metrics=root_mean_squared_error, wd=1e-2, loss_func=torch.nn.SmoothL1Loss())
+    data.num_workers = 0
+    data.bs = 1
+    learner = unet_learner(data, resnet34, metrics=rmse, wd=1e-2, n_out=1, loss_func=torch.nn.SmoothL1Loss())
     learner.fit_one_cycle(1, 1e-3)
 
 
@@ -28,7 +29,7 @@ if __name__ == "__main__":
         print("usage: %s <data_path> <out_folder>" % sys.argv[0], file=sys.stderr)
         sys.exit(0)
 
-    data = create_databunch(sys.argv[1])
+    data = create_data(Path(sys.argv[1]))
     data.batch_size = 1
     data.num_workers = 0
     learner = train(data)