Fast, portable neural networks with Gluon HybridBlocks

In [1]:

    

import mxnet as mx
from mxnet import nd
from mxnet import sym
from mxnet.gluon import nn

from time import time


    

In [2]:

    

def get_net():
    # Construct a MLP
    net = nn.HybridSequential()
    with net.name_scope():
        net.add(nn.Dense(256, activation="relu"))
        net.add(nn.Dense(128, activation="relu"))
        net.add(nn.Dense(2))
    # initialize the parameters
    net.collect_params().initialize()
    return net


    

In [3]:

    

x = nd.random_normal(shape=(1, 512))
net = get_net()
print('=== net(x) ==={}'.format(net(x)))


    

    




=== net(x) ===
[[0.08811305 0.06387275]]
<NDArray 1x2 @cpu(0)>

In [4]:

    

net.hybridize()
print('=== net(x) ==={}'.format(net(x)))


    

    




=== net(x) ===
[[0.08811305 0.06387275]]
<NDArray 1x2 @cpu(0)>

In [5]:

    

def bench(net, x):
    mx.nd.waitall()
    start = time()
    for i in range(1000):
        y = net(x)
    mx.nd.waitall()
    return time() - start

net = get_net()
print('Before hybridizing: %.4f sec' % (bench(net, x)))
net.hybridize()
print('After hybridizing: %.4f sec' % (bench(net, x)))


    

    




Before hybridizing: 0.6067 sec
After hybridizing: 0.3179 sec

In [6]:

    

x = sym.var('data')
print('=== input data holder ===')
print(x)

y = net(x)
print('\n=== the symbolic program of net===')
print(y)

y_json = y.tojson()
print('\n=== the according json definition===')
print(y_json)


    

    




=== input data holder ===
<Symbol data>

=== the symbolic program of net===
<Symbol hybridsequential1_dense2_fwd>

=== the according json definition===
{
  "nodes": [
    {
      "op": "null", 
      "name": "data", 
      "inputs": []
    }, 
    {
      "op": "null", 
      "name": "hybridsequential1_dense0_weight", 
      "attrs": {
        "__dtype__": "0", 
        "__lr_mult__": "1.0", 
        "__shape__": "(256, 0)", 
        "__storage_type__": "0", 
        "__wd_mult__": "1.0"
      }, 
      "inputs": []
    }, 
    {
      "op": "null", 
      "name": "hybridsequential1_dense0_bias", 
      "attrs": {
        "__dtype__": "0", 
        "__init__": "zeros", 
        "__lr_mult__": "1.0", 
        "__shape__": "(256,)", 
        "__storage_type__": "0", 
        "__wd_mult__": "1.0"
      }, 
      "inputs": []
    }, 
    {
      "op": "FullyConnected", 
      "name": "hybridsequential1_dense0_fwd", 
      "attrs": {
        "flatten": "True", 
        "no_bias": "False", 
        "num_hidden": "256"
      }, 
      "inputs": [[0, 0, 0], [1, 0, 0], [2, 0, 0]]
    }, 
    {
      "op": "Activation", 
      "name": "hybridsequential1_dense0_relu_fwd", 
      "attrs": {"act_type": "relu"}, 
      "inputs": [[3, 0, 0]]
    }, 
    {
      "op": "null", 
      "name": "hybridsequential1_dense1_weight", 
      "attrs": {
        "__dtype__": "0", 
        "__lr_mult__": "1.0", 
        "__shape__": "(128, 0)", 
        "__storage_type__": "0", 
        "__wd_mult__": "1.0"
      }, 
      "inputs": []
    }, 
    {
      "op": "null", 
      "name": "hybridsequential1_dense1_bias", 
      "attrs": {
        "__dtype__": "0", 
        "__init__": "zeros", 
        "__lr_mult__": "1.0", 
        "__shape__": "(128,)", 
        "__storage_type__": "0", 
        "__wd_mult__": "1.0"
      }, 
      "inputs": []
    }, 
    {
      "op": "FullyConnected", 
      "name": "hybridsequential1_dense1_fwd", 
      "attrs": {
        "flatten": "True", 
        "no_bias": "False", 
        "num_hidden": "128"
      }, 
      "inputs": [[4, 0, 0], [5, 0, 0], [6, 0, 0]]
    }, 
    {
      "op": "Activation", 
      "name": "hybridsequential1_dense1_relu_fwd", 
      "attrs": {"act_type": "relu"}, 
      "inputs": [[7, 0, 0]]
    }, 
    {
      "op": "null", 
      "name": "hybridsequential1_dense2_weight", 
      "attrs": {
        "__dtype__": "0", 
        "__lr_mult__": "1.0", 
        "__shape__": "(2, 0)", 
        "__storage_type__": "0", 
        "__wd_mult__": "1.0"
      }, 
      "inputs": []
    }, 
    {
      "op": "null", 
      "name": "hybridsequential1_dense2_bias", 
      "attrs": {
        "__dtype__": "0", 
        "__init__": "zeros", 
        "__lr_mult__": "1.0", 
        "__shape__": "(2,)", 
        "__storage_type__": "0", 
        "__wd_mult__": "1.0"
      }, 
      "inputs": []
    }, 
    {
      "op": "FullyConnected", 
      "name": "hybridsequential1_dense2_fwd", 
      "attrs": {
        "flatten": "True", 
        "no_bias": "False", 
        "num_hidden": "2"
      }, 
      "inputs": [[8, 0, 0], [9, 0, 0], [10, 0, 0]]
    }
  ], 
  "arg_nodes": [0, 1, 2, 5, 6, 9, 10], 
  "node_row_ptr": [
    0, 
    1, 
    2, 
    3, 
    4, 
    5, 
    6, 
    7, 
    8, 
    9, 
    10, 
    11, 
    12
  ], 
  "heads": [[11, 0, 0]], 
  "attrs": {"mxnet_version": ["int", 10300]}
}

In [7]:

    

net.export('./models/my_model', epoch=0)