Merge pull request #1 from dronedeploy/update_callbacks

fix callbacks, add image logging

Author: nickponline

libs/training.py

@@ -43,9 +43,10 @@ def train_model(dataset):
 
     data = datasets.load_dataset(dataset, size, bs)
     encoder_model = models.resnet18
-    learn = unet_learner(data, encoder_model, path='models', metrics=metrics, wd=wd, bottle=True, pretrained=pretrained, callback_fns=WandbCallback)
+    learn = unet_learner(data, encoder_model, path='models', metrics=metrics, wd=wd, bottle=True, pretrained=pretrained)
 
     callbacks = [
+        WandbCallback(learn, log=None, input_type="images"),
         MyCSVLogger(learn, filename='baseline_model'),
         ExportCallback(learn, "baseline_model", monitor='f_beta'),
         MySaveModelCallback(learn, every='epoch', monitor='f_beta')

libs/util.py

@@ -1,5 +1,6 @@
 from typing import Any
-from fastai.callbacks import CSVLogger, Callback, SaveModelCallback, TrackerCallback
+from fastai.callbacks import CSVLogger, SaveModelCallback, TrackerCallback
+from fastai.callback import Callback
 from fastai.metrics import add_metrics
 from fastai.torch_core import dataclass, torch, Tensor, Optional, warn
 from fastai.basic_train import Learner
@@ -22,6 +23,7 @@ def on_epoch_end(self, epoch:int, **kwargs:Any)->None:
             self.best = current
             self.learn.export(self.model_path)
 
+# TODO: does this delete some other path or just overwrite?
 class MySaveModelCallback(SaveModelCallback):
     """Saves the model after each epoch to potentially resume training.
 
@@ -48,6 +50,7 @@ def __init__(self, learn, filename='history'):
 
     def on_train_begin(self, **kwargs):
         if self.path.exists():
+            # TODO: does this open a file named "a"...?
             self.file = self.path.open('a')
         else:
             super().on_train_begin(**kwargs)

requirements.txt

@@ -1,5 +1,7 @@
-numpy==1.17.1
-torch==1.1.0
+fastai
 opencv_python==3.4.3.18
+numpy==1.17.1
+sklearn
+torch
 typing==3.6.6
-fastai==1.0.54
+wandb