linearwrap for several examples

examples/Atari2600/DQN.py

@@ -71,21 +71,22 @@ class Model(ModelDesc):
         """ image: [0,255]"""
         image = image / 255.0
         with argscope(Conv2D, nl=PReLU.f, use_bias=True):
-            l = Conv2D('conv0', image, out_channel=32, kernel_shape=5, stride=1)
-            l = MaxPooling('pool0', l, 2)
-            l = Conv2D('conv1', l, out_channel=32, kernel_shape=5, stride=1)
-            l = MaxPooling('pool1', l, 2)
-            l = Conv2D('conv2', l, out_channel=64, kernel_shape=4)
-            l = MaxPooling('pool2', l, 2)
-            l = Conv2D('conv3', l, out_channel=64, kernel_shape=3)
-
-            # the original arch
-            #l = Conv2D('conv0', image, out_channel=32, kernel_shape=8, stride=4)
-            #l = Conv2D('conv1', l, out_channel=64, kernel_shape=4, stride=2)
-            #l = Conv2D('conv2', l, out_channel=64, kernel_shape=3)
-
-        l = FullyConnected('fc0', l, 512, nl=lambda x, name: LeakyReLU.f(x, 0.01, name))
-        l = FullyConnected('fct', l, out_dim=NUM_ACTIONS, nl=tf.identity)
+            l = (LinearWrap(image)
+                .Conv2D('conv0', out_channel=32, kernel_shape=5)
+                .MaxPooling('pool0', 2)
+                .Conv2D('conv1', out_channel=32, kernel_shape=5)
+                .MaxPooling('pool1', 2)
+                .Conv2D('conv2', out_channel=64, kernel_shape=4)
+                .MaxPooling('pool2', 2)
+                .Conv2D('conv3', out_channel=64, kernel_shape=3)
+
+                # the original arch
+                #.Conv2D('conv0', image, out_channel=32, kernel_shape=8, stride=4)
+                #.Conv2D('conv1', out_channel=64, kernel_shape=4, stride=2)
+                #.Conv2D('conv2', out_channel=64, kernel_shape=3)
+
+                .FullyConnected('fc0', 512, nl=lambda x, name: LeakyReLU.f(x, 0.01, name))
+                .FullyConnected('fct', NUM_ACTIONS, nl=tf.identity)())
         return l
 
     def _build_graph(self, inputs, is_training):

examples/DoReFa-Net/svhn-digit-dorefa.py

@@ -101,40 +101,43 @@ class Model(ModelDesc):
         def activate(x):
             return fa(cabs(x))
 
-        l = image / 256.0
+        image = image / 256.0
 
         with argscope(BatchNorm, decay=0.9, epsilon=1e-4, use_local_stat=is_training), \
                 argscope(Conv2D, use_bias=False, nl=tf.identity):
-            l = Conv2D('conv0', l, 48, 5, padding='VALID', use_bias=True)
-            l = MaxPooling('pool0', l, 2, padding='SAME')
-            l = activate(l)
-            # 18
-
-            l = Conv2D('conv1', l, 64, 3, padding='SAME')
-            l = activate(BatchNorm('bn1', fg(l)))
-
-            l = Conv2D('conv2', l, 64, 3, padding='SAME')
-            l = BatchNorm('bn2', fg(l))
-            l = MaxPooling('pool1', l, 2, padding='SAME')
-            l = activate(l)
-            # 9
-            l = Conv2D('conv3', l, 128, 3, padding='VALID')
-            l = activate(BatchNorm('bn3', fg(l)))
-            # 7
-
-            l = Conv2D('conv4', l, 128, 3, padding='SAME')
-            l = activate(BatchNorm('bn4', fg(l)))
-
-            l = Conv2D('conv5', l, 128, 3, padding='VALID')
-            l = activate(BatchNorm('bn5', fg(l)))
-            # 5
-
-            l = tf.nn.dropout(l, 0.5 if is_training else 1.0)
-            l = Conv2D('conv6', l, 512, 5, padding='VALID')
-            l = BatchNorm('bn6', fg(l))
-            l = cabs(l)
-
-            logits = FullyConnected('fc1', l, 10, nl=tf.identity)
+            logits = (LinearWrap(image)
+                .Conv2D('conv0', 48, 5, padding='VALID', use_bias=True)
+                .MaxPooling('pool0', 2, padding='SAME')
+                .apply(activate)
+                # 18
+                .Conv2D('conv1', 64, 3, padding='SAME')
+                .apply(fg)
+                .BatchNorm('bn1').apply(activate)
+
+                .Conv2D('conv2', 64, 3, padding='SAME')
+                .apply(fg)
+                .BatchNorm('bn2')
+                .MaxPooling('pool1', 2, padding='SAME')
+                .apply(activate)
+                # 9
+                .Conv2D('conv3', 128, 3, padding='VALID')
+                .apply(fg)
+                .BatchNorm('bn3').apply(activate)
+                # 7
+
+                .Conv2D('conv4', 128, 3, padding='SAME')
+                .apply(fg)
+                .BatchNorm('bn4').apply(activate)
+
+                .Conv2D('conv5', 128, 3, padding='VALID')
+                .apply(fg)
+                .BatchNorm('bn5').apply(activate)
+                # 5
+                .tf.nn.dropout(0.5 if is_training else 1.0)
+                .Conv2D('conv6', 512, 5, padding='VALID')
+                .apply(fg).BatchNorm('bn6')
+                .apply(cabs)
+                .FullyConnected('fc1', 10, nl=tf.identity)())
         prob = tf.nn.softmax(logits, name='output')
 
         cost = tf.nn.sparse_softmax_cross_entropy_with_logits(logits, label)

examples/load-alexnet.py

@@ -3,22 +3,14 @@
 # File: load-alexnet.py
 # Author: Yuxin Wu <ppwwyyxx@gmail.com>
 
-import cv2  # tf bug
 import tensorflow as tf
 import numpy as np
 import os
 import argparse
-import cPickle as pkl
 
-from tensorpack.train import TrainConfig
-from tensorpack.predict import PredictConfig, get_predict_func
-from tensorpack.models import *
-from tensorpack.utils import *
-from tensorpack.tfutils import *
+from tensorpack import *
 from tensorpack.tfutils.symbolic_functions import *
 from tensorpack.tfutils.summary import *
-from tensorpack.callbacks import *
-from tensorpack.dataflow import *
 from tensorpack.dataflow.dataset import ILSVRCMeta
 
 """

examples/load-vgg16.py

@@ -40,39 +40,35 @@ class Model(ModelDesc):
 
         with argscope(Conv2D, kernel_shape=3):
             # 224
-            l = Conv2D('conv1_1', image, 64)
-            l = Conv2D('conv1_2', l, 64)
-            l = MaxPooling('pool1', l, 2)
-            # 112
-
-            l = Conv2D('conv2_1', l, 128)
-            l = Conv2D('conv2_2', l, 128)
-            l = MaxPooling('pool2', l, 2)
-            # 56
-
-            l = Conv2D('conv3_1', l, 256)
-            l = Conv2D('conv3_2', l, 256)
-            l = Conv2D('conv3_3', l, 256)
-            l = MaxPooling('pool3', l, 2)
-            # 28
-
-            l = Conv2D('conv4_1', l, 512)
-            l = Conv2D('conv4_2', l, 512)
-            l = Conv2D('conv4_3', l, 512)
-            l = MaxPooling('pool4', l, 2)
-            # 14
-
-            l = Conv2D('conv5_1', l, 512)
-            l = Conv2D('conv5_2', l, 512)
-            l = Conv2D('conv5_3', l, 512)
-            l = MaxPooling('pool5', l, 2)
-        # 7
-
-        l = FullyConnected('fc6', l, 4096)
-        l = tf.nn.dropout(l, keep_prob)
-        l = FullyConnected('fc7', l, 4096)
-        l = tf.nn.dropout(l, keep_prob)
-        logits = FullyConnected('fc8', l, out_dim=1000, nl=tf.identity)
+            logits = (LinearWrap(image)
+                .Conv2D('conv1_1', 64)
+                .Conv2D('conv1_2', 64)
+                .MaxPooling('pool1', 2)
+                # 112
+                .Conv2D('conv2_1', 128)
+                .Conv2D('conv2_2', 128)
+                .MaxPooling('pool2', 2)
+                # 56
+                .Conv2D('conv3_1', 256)
+                .Conv2D('conv3_2', 256)
+                .Conv2D('conv3_3', 256)
+                .MaxPooling('pool3', 2)
+                # 28
+                .Conv2D('conv4_1', 512)
+                .Conv2D('conv4_2', 512)
+                .Conv2D('conv4_3', 512)
+                .MaxPooling('pool4', 2)
+                # 14
+                .Conv2D('conv5_1', 512)
+                .Conv2D('conv5_2', 512)
+                .Conv2D('conv5_3', 512)
+                .MaxPooling('pool5', 2)
+                 # 7
+                .FullyConnected('fc6', 4096)
+                .tf.nn.dropout(keep_prob)
+                .FullyConnected('fc7', 4096)
+                .tf.nn.dropout(keep_prob)
+                .FullyConnected('fc8', out_dim=1000, nl=tf.identity)())
         prob = tf.nn.softmax(logits, name='output')
 
 def run_test(path, input):
@@ -104,8 +100,8 @@ def run_test(path, input):
 
 if __name__ == '__main__':
     parser = argparse.ArgumentParser()
-    parser.add_argument('--gpu', default='0',
-                        help='comma separated list of GPU(s) to use.') # nargs='*' in multi mode
+    parser.add_argument('--gpu',
+                        help='comma separated list of GPU(s) to use.')
     parser.add_argument('--load', required=True,
                         help='.npy model file generated by tensorpack.utils.loadcaffe')
     parser.add_argument('--input', help='an input image', required=True)

examples/mnist-convnet.py

@@ -34,16 +34,16 @@ class Model(ModelDesc):
         nl = PReLU.f
         image = image * 2 - 1
         with argscope(Conv2D, kernel_shape=3, nl=nl, out_channel=32):
-            logits = LinearWrap(image) \
-                    .Conv2D('conv0', padding='VALID') \
-                    .MaxPooling('pool0', 2) \
-                    .Conv2D('conv1', padding='SAME') \
-                    .Conv2D('conv2', padding='VALID') \
-                    .MaxPooling('pool1', 2) \
-                    .Conv2D('conv3', padding='VALID') \
-                    .FullyConnected('fc0', 512) \
-                    .tf.nn.dropout(keep_prob) \
-                    .FullyConnected('fc1', out_dim=10, nl=tf.identity)()
+            logits = (LinearWrap(image) # the starting brace is only for line-breaking
+                    .Conv2D('conv0', padding='VALID')
+                    .MaxPooling('pool0', 2)
+                    .Conv2D('conv1', padding='SAME')
+                    .Conv2D('conv2', padding='VALID')
+                    .MaxPooling('pool1', 2)
+                    .Conv2D('conv3', padding='VALID')
+                    .FullyConnected('fc0', 512)
+                    .tf.nn.dropout(keep_prob)
+                    .FullyConnected('fc1', out_dim=10, nl=tf.identity)())
         prob = tf.nn.softmax(logits, name='prob')
 
         cost = tf.nn.sparse_softmax_cross_entropy_with_logits(logits, label)

examples/svhn-digit-convnet.py

@@ -28,17 +28,17 @@ class Model(ModelDesc):
 
         image = image / 128.0 - 1
 
-        l = Conv2D('conv1', image, 24, 5, padding='VALID')
-        l = MaxPooling('pool1', l, 2, padding='SAME')
-        l = Conv2D('conv2', l, 32, 3, padding='VALID')
-        l = Conv2D('conv3', l, 32, 3, padding='VALID')
-        l = MaxPooling('pool2', l, 2, padding='SAME')
-        l = Conv2D('conv4', l, 64, 3, padding='VALID')
-
-        l = tf.nn.dropout(l, keep_prob)
-        l = FullyConnected('fc0', l, 512,
-                           b_init=tf.constant_initializer(0.1))
-        logits = FullyConnected('linear', l, out_dim=10, nl=tf.identity)
+        logits = LinearWrap(image) \
+                .Conv2D('conv1', 24, 5, padding='VALID') \
+                .MaxPooling('pool1', 2, padding='SAME') \
+                .Conv2D('conv2', 32, 3, padding='VALID') \
+                .Conv2D('conv3', 32, 3, padding='VALID') \
+                .MaxPooling('pool2', 2, padding='SAME') \
+                .Conv2D('conv4', 64, 3, padding='VALID') \
+                .tf.nn.dropout(keep_prob) \
+                .FullyConnected('fc0', 512,
+                        b_init=tf.constant_initializer(0.1)) \
+                .FullyConnected('linear', out_dim=10, nl=tf.identity)()
         prob = tf.nn.softmax(logits, name='output')
 
         cost = tf.nn.sparse_softmax_cross_entropy_with_logits(logits, label)

tensorpack/models/__init__.py

@@ -54,10 +54,15 @@ class LinearWrap(object):
             return f
         else:
             if layer_name != 'tf':
-                logger.warn("You're calling LinearWrap with something neither a layer nor 'tf'. Not officially supported yet!")
+                logger.warn("You're calling LinearWrap.__getattr__ with something neither a layer nor 'tf'. Not officially supported yet!")
             assert isinstance(layer, ModuleType)
             return LinearWrap.TFModuleFunc(layer, self._t)
 
+    def apply(self, func, *args, **kwargs):
+        """ send tensor to the first argument of a simple func"""
+        ret = func(self._t, *args, **kwargs)
+        return LinearWrap(ret)
+
     def __call__(self):
         return self._t