update for ilsvrc

scripts/dump_train_config.py

@@ -10,7 +10,7 @@
 import tqdm
 import os
 from tensorpack.utils import logger
-from tensorpack.utils.utils import mkdir_p
+from tensorpack.utils.fs import mkdir_p
 
 
 parser = argparse.ArgumentParser()

tensorpack/dataflow/dataset/ilsvrc.py

@@ -4,8 +4,8 @@
 # Author: Yuxin Wu <ppwwyyxxc@gmail.com>
 import os
 import tarfile
+import cv2
 import numpy as np
-import scipy.ndimage as scimg
 
 from ...utils import logger, get_rng
 from ..base import DataFlow
@@ -61,9 +61,10 @@ def get_image_list(self, name):
                 ret.append((name, int(cls)))
             return ret
 
-    def get_per_pixel_mean(self):
+    def get_per_pixel_mean(self, size=None):
         """
-        :returns per-pixel mean as an array of shape (3, 256, 256) in range [0, 255]
+        :param size: return image size in [h, w]. default to (256, 256)
+        :returns per-pixel mean as an array of shape (h, w, 3) in range [0, 255]
         """
         import imp
         caffepb = imp.load_source('caffepb', self.caffe_pb_file)
@@ -73,6 +74,9 @@ def get_per_pixel_mean(self):
         with open(mean_file) as f:
             obj.ParseFromString(f.read())
         arr = np.array(obj.data).reshape((3, 256, 256))
+        arr = np.transpose(arr, [1,2,0])
+        if size is not None:
+            arr = cv2.resize(arr, size[::-1])
         return arr
 
 class ILSVRC12(DataFlow):
@@ -106,9 +110,10 @@ def get_data(self):
             self.rng.shuffle(idxs)
         for k in idxs:
             tp = self.imglist[k]
-            fname = os.path.join(self.dir, self.name, tp[0])
-            im = scimg.imread(fname)
-            if len(im.shape) == 2:
+            fname = os.path.join(self.dir, self.name, tp[0]).strip()
+            im = cv2.imread(fname)
+            assert im is not None, fname
+            if im.ndim == 2:
                 im = np.expand_dims(im, 2).repeat(3,2)
             yield [im, tp[1]]
 

tensorpack/dataflow/imgaug/base.py

@@ -69,7 +69,7 @@ def __init__(self, augmentors):
         self.augs = augmentors
 
     def _augment(self, img):
-        assert img.arr.ndim in [2, 3]
+        assert img.arr.ndim in [2, 3], img.arr.ndim
         img.arr = img.arr.astype('float32')
         for aug in self.augs:
             aug.augment(img)

tensorpack/train/trainer.py

@@ -86,11 +86,12 @@ class QueueInputTrainer(Trainer):
     @staticmethod
     def _average_grads(tower_grads):
         ret = []
-        for grad_and_vars in zip(*tower_grads):
-            grad = tf.add_n([x[0] for x in grad_and_vars]) / float(len(tower_grads))
-            v = grad_and_vars[0][1]
-            ret.append((grad, v))
-        return ret
+        with tf.device('/gpu:0'):
+            for grad_and_vars in zip(*tower_grads):
+                grad = tf.add_n([x[0] for x in grad_and_vars]) / float(len(tower_grads))
+                v = grad_and_vars[0][1]
+                ret.append((grad, v))
+            return ret
 
     def train(self):
         model = self.model
@@ -121,7 +122,8 @@ def get_model_inputs():
                     if i == 0:
                         cost_var_t0 = cost_var
                     grad_list.append(
-                        self.config.optimizer.compute_gradients(cost_var))
+                        self.config.optimizer.compute_gradients(cost_var,
+                                                                gate_gradients=0))
 
                     if i == 0:
                         tf.get_variable_scope().reuse_variables()