superres

miniai/augment.py

@@ -51,18 +51,22 @@ def show_image_batch(self:Learner, max_n=9, cbs=None, **kwargs):
 
 # %% ../nbs/14_augment.ipynb 38
 class CapturePreds(Callback):
-    def before_fit(self, learn): self.all_preds,self.all_targs = [],[]
+    def before_fit(self, learn): self.all_inps,self.all_preds,self.all_targs = [],[],[]
     def after_batch(self, learn):
+        self.all_inps. append(to_cpu(learn.batch[0]))
         self.all_preds.append(to_cpu(learn.preds))
         self.all_targs.append(to_cpu(learn.batch[1]))
-    def after_fit(self, learn): self.all_preds,self.all_targs = torch.cat(self.all_preds),torch.cat(self.all_targs)
+    def after_fit(self, learn):
+        self.all_preds,self.all_targs,self.all_inps = map(torch.cat, [self.all_preds,self.all_targs,self.all_inps])
 
 # %% ../nbs/14_augment.ipynb 39
 @fc.patch
-def capture_preds(self: Learner, cbs=None):
+def capture_preds(self: Learner, cbs=None, inps=False):
     cp = CapturePreds()
     self.fit(1, train=False, cbs=[cp]+fc.L(cbs))
-    return cp.all_preds,cp.all_targs
+    res = cp.all_preds,cp.all_targs
+    if inps: res = res+(cp.all_inps,)
+    return res
 
 # %% ../nbs/14_augment.ipynb 54
 def _rand_erase1(x, pct, xm, xs, mn, mx):

miniai/conv.py

@@ -4,6 +4,7 @@
 __all__ = ['def_device', 'conv', 'to_device', 'collate_device']
 
 # %% ../nbs/07_convolutions.ipynb 2
+import torch
 from torch import nn
 
 from torch.utils.data import default_collate

nbs/07_convolutions.ipynb

@@ -23,6 +23,7 @@
    "outputs": [],
    "source": [
     "#|export\n",
+    "import torch\n",
     "from torch import nn\n",
     "\n",
     "from torch.utils.data import default_collate\n",

nbs/14_augment.ipynb

@@ -1418,11 +1418,13 @@
    "source": [
     "#| export\n",
     "class CapturePreds(Callback):\n",
-    "    def before_fit(self, learn): self.all_preds,self.all_targs = [],[]\n",
+    "    def before_fit(self, learn): self.all_inps,self.all_preds,self.all_targs = [],[],[]\n",
     "    def after_batch(self, learn):\n",
+    "        self.all_inps. append(to_cpu(learn.batch[0]))\n",
     "        self.all_preds.append(to_cpu(learn.preds))\n",
     "        self.all_targs.append(to_cpu(learn.batch[1]))\n",
-    "    def after_fit(self, learn): self.all_preds,self.all_targs = torch.cat(self.all_preds),torch.cat(self.all_targs)"
+    "    def after_fit(self, learn):\n",
+    "        self.all_preds,self.all_targs,self.all_inps = map(torch.cat, [self.all_preds,self.all_targs,self.all_inps])"
    ]
   },
   {
@@ -1433,10 +1435,12 @@
    "source": [
     "#| export\n",
     "@fc.patch\n",
-    "def capture_preds(self: Learner, cbs=None):\n",
+    "def capture_preds(self: Learner, cbs=None, inps=False):\n",
     "    cp = CapturePreds()\n",
     "    self.fit(1, train=False, cbs=[cp]+fc.L(cbs))\n",
-    "    return cp.all_preds,cp.all_targs"
+    "    res = cp.all_preds,cp.all_targs\n",
+    "    if inps: res = res+(cp.all_inps,)\n",
+    "    return res"
    ]
   },
   {