From 735ad4f94c591ad9935b83fb6549f3c11c3128b0 Mon Sep 17 00:00:00 2001
From: Bolei Zhou <bzhou@visiongpu15.csail.mit.edu>
Date: Fri, 29 Sep 2017 14:19:32 -0400
Subject: [PATCH] fixed the netdissect pytorch issues of image directory and
 input model

---
 script/rundissect_pytorch_external.sh |  1 -
 src/netprobe_pytorch.py               | 15 +++++----------
 src/viewprobe.py                      |  4 ++--
 3 files changed, 7 insertions(+), 13 deletions(-)

diff --git a/script/rundissect_pytorch_external.sh b/script/rundissect_pytorch_external.sh
index 9fbfed3..7a13979 100755
--- a/script/rundissect_pytorch_external.sh
+++ b/script/rundissect_pytorch_external.sh
@@ -4,7 +4,6 @@ DIR=pytorch_alexnet_imagenet
 ARCH='alexnet' # [alexnet,squeezenet1_1,resnet18,...]. It should work for all the models in https://github.com/pytorch/vision/tree/master/torchvision/models
 LAYERS="features"
 DATASET=dataset/broden1_224
-WEIGHTS="none"
 NUMCLASSES=1000
 
 # default setting
diff --git a/src/netprobe_pytorch.py b/src/netprobe_pytorch.py
index 03d63a9..a8c5c62 100644
--- a/src/netprobe_pytorch.py
+++ b/src/netprobe_pytorch.py
@@ -16,9 +16,6 @@
 import sys
 
 os.environ['GLOG_minloglevel'] = '2'
-import caffe
-from caffe.proto import caffe_pb2
-from google.protobuf import text_format
 from scipy.misc import imresize, imread
 from scipy.ndimage.filters import gaussian_filter
 from scipy.ndimage.interpolation import zoom
@@ -29,8 +26,6 @@
 import rotate
 import expdir
 
-caffe.set_mode_gpu()
-caffe.set_device(0)
 
 def create_probe(
         directory, dataset, definition, weights, mean, blobs,
@@ -47,8 +42,8 @@ def create_probe(
     '''
     directory: where to place the probe_conv5.mmap files.
     data: the AbstractSegmentation data source to draw upon
-    definition: the filename for the caffe prototxt
-    weights: the filename for the caffe model weights
+    definition: the filename for the pytorch
+    weights: the filename for the weights
     mean: to use to normalize rgb values for the network
     blobs: ['conv3', 'conv4', 'conv5'] to probe
     '''
@@ -57,7 +52,7 @@ def create_probe(
     data = loadseg.SegmentationData(args.dataset)
 
     # the network to dissect
-    if args.weights == None:
+    if args.weights == 'none':
         # load the imagenet pretrained model
         net = torchvision.models.__dict__[args.definition](pretrained=True)
     else:
@@ -158,7 +153,7 @@ def hook_feature(module, input, output):
         # previous feedforward case
         inp = inp[:,::-1,:,:]
         inp_tensor = V(torch.from_numpy(inp.copy()))
-        inp_tensor.div_(255.0*0.224) # approximately normalize the input to make the images scaled at around 1.
+        inp_tensor.div_(255.0*0.224) # hack: approximately normalize the input to make the images scaled at around 1.
         inp_tensor = inp_tensor.cuda()
         result = net.forward(inp_tensor)
         # output the hooked feature
@@ -227,7 +222,7 @@ def report(txt):
         import loadseg
 
         parser = argparse.ArgumentParser(description=
-            'Probe a caffe network and save results in a directory.')
+            'Probe a pytorch network and save results in a directory.')
         parser.add_argument(
                 '--directory',
                 default='.',
diff --git a/src/viewprobe.py b/src/viewprobe.py
index 0e48b05..eae4736 100644
--- a/src/viewprobe.py
+++ b/src/viewprobe.py
@@ -109,13 +109,13 @@ def generate_html_summary(self, layer,
                       col*(imsize+1):col*(imsize+1)+imsize,:] = vis
             imfn = 'image/%s%s-%04d.jpg' % (
                     expdir.fn_safe(layer), gridname, unit)
-            imsave(self.ed.filename(['html', imfn]), tiled)
+            imsave(self.ed.filename(os.path.join('html', imfn)), tiled)
             labels = '; '.join(['%s (%s, %f)' %
                     (name_pciou[c][unit], categories[c], score_pciou[c, unit])
                     for c in bestcat_pciou[:,unit]])
             html.extend([
                 '<h6>%s unit %d: %s</h6>' % (layer, unit + 1, labels),
-                '<img src="%s" height="%d">' % (imfn, imscale)
+                '<img src="%s" height="%d">' % (os.path.join('html', imfn), imscale)
                 ])
         html.extend([
             '</div>', '</body>', '</html>', ''])