fixed soft inference + nbdt script

Author: alvinwan

examples/load_pretrained_nbdts.ipynb

@@ -0,0 +1,114 @@
+{
+ "cells": [
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "%pip install nbdt"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "from nbdt.model import SoftNBDT\n",
+    "from nbdt.models import ResNet18, wrn28_10_cifar10, wrn28_10_cifar100, wrn28_10  # use wrn28_10 for TinyImagenet200\n",
+    "from torchvision import transforms\n",
+    "from nbdt.utils import DATASET_TO_CLASSES, load_image_from_path, maybe_install_wordnet\n",
+    "from IPython.display import display"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "model = wrn28_10_cifar10()\n",
+    "model = SoftNBDT(\n",
+    "  pretrained=True,\n",
+    "  dataset='CIFAR10',\n",
+    "  arch='wrn28_10_cifar10',\n",
+    "  model=model)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "image_urls = {\n",
+    "    'cat': 'https://images.pexels.com/photos/126407/pexels-photo-126407.jpeg?auto=compress&cs=tinysrgb&dpr=2&w=300',\n",
+    "    'bear': 'https://images.pexels.com/photos/1466592/pexels-photo-1466592.jpeg?auto=compress&cs=tinysrgb&dpr=2&w=300',\n",
+    "    'dog': 'https://images.pexels.com/photos/1490908/pexels-photo-1490908.jpeg?auto=compress&cs=tinysrgb&dpr=2&w=300'\n",
+    "}"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# show image\n",
+    "im = load_image_from_path(image_urls['cat'])\n",
+    "display(im)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# load + transform image\n",
+    "transforms = transforms.Compose([\n",
+    "  transforms.Resize(32),\n",
+    "  transforms.CenterCrop(32),\n",
+    "  transforms.ToTensor(),\n",
+    "  transforms.Normalize((0.4914, 0.4822, 0.4465), (0.2023, 0.1994, 0.2010)),\n",
+    "])\n",
+    "x = transforms(im)[None]"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# run inference\n",
+    "outputs = model(x)  # to get intermediate decisions, use `model.forward_with_decisions(x)` and add `hierarchy='wordnet' to SoftNBDT\n",
+    "_, predicted = outputs.max(1)\n",
+    "cls = DATASET_TO_CLASSES['CIFAR10'][predicted[0]]\n",
+    "print(cls)"
+   ]
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "pytorch-1.2",
+   "language": "python",
+   "name": "pytorch-1.2"
+  },
+  "language_info": {
+   "codemirror_mode": {
+    "name": "ipython",
+    "version": 3
+   },
+   "file_extension": ".py",
+   "mimetype": "text/x-python",
+   "name": "python",
+   "nbconvert_exporter": "python",
+   "pygments_lexer": "ipython3",
+   "version": "3.7.4"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 4
+}

nbdt/analysis.py

@@ -242,7 +242,8 @@ class SoftEmbeddedDecisionRules(HardEmbeddedDecisionRules):
     name = 'NBDT-Soft'
 
     def forward_with_decisions(self, outputs):
-        predicted = self.forward(outputs)
+        outputs = self.forward(outputs)
+        _, predicted = outputs.max(1)
 
         decisions = []
         node = self.nodes[0]
@@ -251,7 +252,7 @@ def forward_with_decisions(self, outputs):
             leaf = node.wnids[prediction]
             decision = leaf_to_path_nodes[leaf]
             decisions.append(decision)
-        return predicted, decisions
+        return outputs, decisions
 
     def forward(self, outputs):
         outputs = SoftTreeSupLoss.inference(

nbdt/bin/nbdt

@@ -32,8 +32,8 @@ x = transforms(im)[None]
 
 # run inference
 outputs, decisions = model.forward_with_decisions(x)  # use `model(x)` to obtain just logits
-outputs = [3]
-cls = DATASET_TO_CLASSES['CIFAR10'][outputs[0]]
+_, predicted = outputs.max(1)
+cls = DATASET_TO_CLASSES['CIFAR10'][predicted[0]]
 print('Prediction:', cls, '// Decisions:', ', '.join([
     info['name'] for info in decisions[0]
 ][1:]))  # [1:] to skip the root