update sym_slice

docs/mrt/gen.md

@@ -86,7 +86,10 @@ $$
 $$
 
 $$
-Xi = \text{slice}\Big(X, \text{begin=(None,)*i+(i,)+(None,)*(ndims-i-1)}, \text{end=(None,)*i+(i+1,)+(None,)*(ndims-i-1)}, \text{step=step}\Big)
+Xi = \text{slice}\Big(X, \\
+\text{begin=(None,)*ichannel+(i,)+(None,)*(ndims-ichannel-1)}, \\
+\text{end=(None,)*ichannel+(i+step,)+(None,)*(ndims-ichannel-1)}, \\
+\text{step=(-1,)*ichannel+(step_size,)+(-1,)*(ndims-ichannel-1)}\Big)
 $$
 
 If $X$ is of channel feature and $W$ is of layer feature or vice versa, $W$ (or $X$) will also be split to be compatible with $X$ (or $W$).

python/mrt/gen/tfm_ops.py

@@ -14,7 +14,8 @@
 from mrt.tfm_pass import OUT_KEY
 from .tfm_base import register_pass, register_transformer, Transformer
 from .tfm_types import get_quantizer, USQuantizer, UAQuantizer, \
-                       FT_TYPE_EXP, AFeature, SBuffer
+                       FT_TYPE_EXP, AFeature, SBuffer, LAYER_WISE_TYPE, \
+                       US_QUANT_TYPE
 from .tfm_utils import scale_exp, get_buffer_exp
 
 from mrt import sim_quant_helper as sim
@@ -119,6 +120,49 @@ def quantize(self, op, **kwargs):
 @register_pass("prepare_for_compile")
 @register_transformer("Convolution")
 class Convolution(tops.Convolution):
+    def slice_channel(self, op, **kwargs):
+        name, op_name = op.attr('name'), op.attr('op_name')
+        attr, childs = op.list_attr(), sym_iter(op.get_children())
+        cns = [c.attr('name') for c in childs]
+        cfg_dict = kwargs['cfg_dict']
+        infer_shapes = kwargs['infer_shapes']
+
+        gn_info = cfg_dict[name]['gn_info']
+        ichannel, step = gn_info['ichannel'], gn_info['step']
+
+        assert len(childs) == 2
+        X, W = childs
+        xshp = infer_shapes[cns[0]][get_entry_id(childs[0])]
+        wshp = infer_shapes[cns[1]][get_entry_id(childs[1])]
+        assert len(xshp) == len(wshp) == 4 and \
+            xshp[1] == wshp[1] and xshp[1]%step == 0
+
+        xi_cfg_info, wi_cfg_info = cfg_dict[cns[0]], cfg_dict[cns[1]]
+        xi_cfg_info['gn_info'] = {'gn_type': LAYER_WISE_TYPE}
+        wi_cfg_info['gn_info'] = {'gn_type': LAYER_WISE_TYPE}
+        yi_cfg_info = {
+            'gn_info': {'gn_type': LAYER_WISE_TYPE},
+            'quant_type': US_QUANT_TYPE,
+            'opt_info': cfg_dict[name]['opt_info'],
+        }
+        xs = sym_slice(X, ichannel, step, **kwargs)
+        ws = sym_slice(W, ichannel, step, **kwargs)
+
+        nodes = []
+        for i in range(0, xshp[1], step):
+            suffix = '_' + str(i)+'-'+str(i+step)
+            xni = xs[i].attr('name')
+            cfg_dict[xni] = xi_cfg_info
+            wni = ws[i].attr('name')
+            cfg_dict[wni] = wi_cfg_info
+            yni = N.n(name+suffix)
+            Yi = get_mxnet_op(op_name)(xs[i], ws[i], **attr, name=yni)
+            cfg_dict[yni] = yi_cfg_info
+            nodes.append(Yi)
+
+        op = mx.sym.add_n(*nodes, name=name)
+        return op
+
     def rewrite(self, op, **kwargs):
         op = super().rewrite(op, **kwargs)
         op = separate_pad(op, **kwargs)
@@ -312,3 +356,20 @@ def _quantize_xw(op, **kwargs):
     infer_prec = kprec + xprec + wprec
     precs[name][OUT_KEY] = infer_prec
     return op
+
+def sym_slice(op, ichannel, step, **kwargs):
+    name = op.attr('name')
+    shp = kwargs['infer_shapes'][name][get_entry_id(op)]
+    ndims = len(shp)
+    nodes = []
+    rchannel = ndims-ichannel-1
+    for i in range(0, shp[ichannel], step):
+        suffix = '_' + str(i)+'-'+str(i+step)
+        opi = mx.sym.slice(
+            op, begin=(None,)*ichannel+(i,)+(None,)*rchannel,
+            end=(None,)*ichannel+(i+step,)+(None,)*rchannel,
+            step=(-1,)*ichannel+(step,)+(-1,)*rchannel,
+            name=N.n(name+suffix))
+        nodes.append(opi)
+    return nodes
+

python/mrt/gen/tfm_pass.py

@@ -18,29 +18,6 @@
 
 from mrt import sym_utils as sutils
 
-#----------------------------
-# Channel Slice interfaces
-#----------------------------
-
-def sym_slice_channel(symbol, params, cfg_dict={}):
-    """ Customized graph-level topo pass definition.
-
-        Interface for granularity control.
-        While layer-wise feature is by default,
-        MRT support channel-wise features specified in cfg_dict.
-    """
-
-    def _slice_channel(op, **kwargs):
-        name, op_name = op.attr("name"), op.attr("op_name")
-        gn_info = cfg_dict[name].get("gn_info", DEFAULT_GN_INFO)
-        gn_type = gn_info["gn_type"]
-        if gn_type == CHANNEL_WISE_TYPE:
-            op = apply_pass(
-                "slice_channel", cfg_dict=cfg_dict)(op, **kwargs)
-        return op
-
-    return topo_visit_transformer(symbol, params, _slice_channel)
-
 #----------------------------
 # Module main interfaces
 #----------------------------
@@ -291,6 +268,32 @@ def rewrite(symbol, params):
     return topo_visit_transformer(symbol, params,
             apply_pass("rewrite", infer_shapes=infer_shapes))
 
+#----------------------------
+# Channel Slice interfaces
+#----------------------------
+
+@N.register_nm("slice_channel")
+def sym_slice_channel(symbol, params, cfg_dict={}):
+    """ Customized graph-level topo pass definition.
+
+        Interface for granularity control.
+        While layer-wise feature is by default,
+        MRT support channel-wise features specified in cfg_dict.
+    """
+    infer_shapes = infer_shape(symbol, params)
+
+    def _slice_channel(op, **kwargs):
+        name, op_name = op.attr("name"), op.attr("op_name")
+        gn_info = cfg_dict[name].get("gn_info", DEFAULT_GN_INFO)
+        gn_type = gn_info["gn_type"]
+        if gn_type == CHANNEL_WISE_TYPE:
+            op = apply_pass(
+                "slice_channel", cfg_dict=cfg_dict,
+                infer_shapes=infer_shapes)(op, **kwargs)
+        return op
+
+    return topo_visit_transformer(symbol, params, _slice_channel)
+
 #----------------------------
 # Module quantize interfaces
 #----------------------------

python/mrt/gen/tfm_types.py

@@ -268,8 +268,10 @@ def int_realize(self, data, prec, **kwargs):
             " base `int_realize` function " + \
             "defined in Quantizer")
 
+US_QUANT_TYPE = "UniformSymmetric"
 
-@register_quantizer("UniformSymmetric")
+
+@register_quantizer(US_QUANT_TYPE)
 class USQuantizer(Quantizer):
     """ Information data type for uniform symmetric quantizaton
     """
@@ -493,7 +495,7 @@ def int_realize(self, data, prec, **kwargs):
         out = out.clip(a_min=lower, a_max=upper)
         return out, self.get_prec(out)
 
-DEFAULT_QUANT_TYPE = "UniformSymmetric"
+DEFAULT_QUANT_TYPE = US_QUANT_TYPE
 DEFAULT_QUANTIZER = USQuantizer()
 
 QUANT_INSTANCES = {