1# Licensed to the Apache Software Foundation (ASF) under one
2# or more contributor license agreements.  See the NOTICE file
3# distributed with this work for additional information
4# regarding copyright ownership.  The ASF licenses this file
5# to you under the Apache License, Version 2.0 (the
6# "License"); you may not use this file except in compliance
7# with the License.  You may obtain a copy of the License at
8#
9#   http://www.apache.org/licenses/LICENSE-2.0
10#
11# Unless required by applicable law or agreed to in writing,
12# software distributed under the License is distributed on an
13# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
14# KIND, either express or implied.  See the License for the
15# specific language governing permissions and limitations
16# under the License.
17"""
18Compile Darknet Models for RNN
19==============================
20**Author**: `Siju Samuel <https://siju-samuel.github.io/>`_
21
22This article is an introductory tutorial to deploy darknet rnn models with NNVM.
23
24This script will run a character prediction model
25Each module consists of 3 fully-connected layers. The input layer propagates information from the
26input to the current state. The recurrent layer propagates information through time from the
27previous state to the current one.
28
29The input to the network is a 1-hot encoding of ASCII characters. We train the network to predict
30the next character in a stream of characters. The output is constrained to be a probability
31distribution using a softmax layer.
32
33Since each recurrent layer contains information about the current character and the past
34characters, it can use this context to predict the future characters in a word or phrase.
35
36All the required models and libraries will be downloaded from the internet
37by the script.
38"""
39import random
40import numpy as np
41import tvm
42from tvm.contrib import graph_runtime
43from tvm.contrib.download import download_testdata
44from nnvm.testing.darknet import __darknetffi__
45import nnvm
46import nnvm.frontend.darknet
47
48# Set the parameters
49# -----------------------
50# Set the seed value and the number of characters to predict
51
52#Model name
53MODEL_NAME = 'rnn'
54#Seed value
55seed = 'Thus'
56#Number of characters to predict
57num = 1000
58
59# Download required files
60# -----------------------
61# Download cfg and weights file if first time.
62CFG_NAME = MODEL_NAME + '.cfg'
63WEIGHTS_NAME = MODEL_NAME + '.weights'
64REPO_URL = 'https://github.com/dmlc/web-data/blob/master/darknet/'
65CFG_URL = REPO_URL + 'cfg/' + CFG_NAME + '?raw=true'
66WEIGHTS_URL = REPO_URL + 'weights/' + WEIGHTS_NAME + '?raw=true'
67
68cfg_path = download_testdata(CFG_URL, CFG_NAME, module='darknet')
69weights_path = download_testdata(WEIGHTS_URL, WEIGHTS_NAME, module='darknet')
70
71# Download and Load darknet library
72DARKNET_LIB = 'libdarknet.so'
73DARKNET_URL = REPO_URL + 'lib/' + DARKNET_LIB + '?raw=true'
74lib_path = download_testdata(DARKNET_URL, DARKNET_LIB, module='darknet')
75DARKNET_LIB = __darknetffi__.dlopen(lib_path)
76net = DARKNET_LIB.load_network(cfg_path.encode('utf-8'), weights_path.encode('utf-8'), 0)
77dtype = 'float32'
78batch_size = 1
79
80# Import the graph to NNVM
81# ------------------------
82# Import darknet graph definition to nnvm.
83#
84# Results:
85#   sym: nnvm graph for rnn model
86#   params: params converted from darknet weights
87print("Converting darknet rnn model to nnvm symbols...")
88sym, params = nnvm.frontend.darknet.from_darknet(net, dtype)
89
90# Compile the model on NNVM
91data = np.empty([1, net.inputs], dtype)#net.inputs
92
93target = 'llvm'
94shape = {'data': data.shape}
95print("Compiling the model...")
96
97shape_dict = {'data': data.shape}
98dtype_dict = {'data': data.dtype}
99
100with nnvm.compiler.build_config(opt_level=2):
101    graph, lib, params = nnvm.compiler.build(sym, target, shape_dict, dtype_dict, params)
102
103# Execute the portable graph on TVM
104# ---------------------------------
105# Now we can try deploying the NNVM compiled model on cpu target.
106
107# Set the cpu context
108ctx = tvm.cpu(0)
109# Create graph runtime
110m = graph_runtime.create(graph, lib, ctx)
111# Set the params to runtime
112m.set_input(**params)
113
114def _init_state_memory(rnn_cells_count, dtype):
115    '''Initialize memory for states'''
116    states = {}
117    state_shape = (1024,)
118    for i in range(rnn_cells_count):
119        k = 'rnn' + str(i) + '_state'
120        states[k] = tvm.nd.array(np.zeros(state_shape, dtype).astype(dtype))
121    return states
122
123def _set_state_input(runtime, states):
124    '''Set the state inputs'''
125    for state in states:
126        runtime.set_input(state, states[state])
127
128def _get_state_output(runtime, states):
129    '''Get the state outputs and save'''
130    i = 1
131    for state in states:
132        data = states[state]
133        states[state] = runtime.get_output((i), tvm.nd.empty(data.shape, data.dtype))
134        i += 1
135
136def _proc_rnn_output(out_data):
137    '''Generate the characters from the output array'''
138    sum_array = 0
139    n = out_data.size
140    r = random.uniform(0, 1)
141    for j in range(n):
142        if out_data[j] < 0.0001:
143            out_data[j] = 0
144        sum_array += out_data[j]
145
146    for j in range(n):
147        out_data[j] *= float(1.0) / sum_array
148        r = r - out_data[j]
149        if r <= 0:
150            return j
151    return n-1
152
153print("RNN generaring text...")
154
155out_shape = (net.outputs,)
156rnn_cells_count = 3
157
158# Initialize state memory
159# -----------------------
160states = _init_state_memory(rnn_cells_count, dtype)
161
162len_seed = len(seed)
163count = len_seed + num
164out_txt = ""
165
166#Initialize random seed
167random.seed(0)
168c = ord(seed[0])
169inp_data = np.zeros([net.inputs], dtype)
170
171# Run the model
172# -------------
173
174# Predict character by character till `num`
175for i in range(count):
176    inp_data[c] = 1
177
178    # Set the input data
179    m.set_input('data', tvm.nd.array(inp_data.astype(dtype)))
180    inp_data[c] = 0
181
182    # Set the state inputs
183    _set_state_input(m, states)
184
185    # Run the model
186    m.run()
187
188    # Get the output
189    tvm_out = m.get_output(0, tvm.nd.empty(out_shape, dtype)).asnumpy()
190
191    # Get the state outputs
192    _get_state_output(m, states)
193
194    # Get the predicted character and keep buffering it
195    c = ord(seed[i])  if i < len_seed else _proc_rnn_output(tvm_out)
196    out_txt += chr(c)
197
198print("Predicted Text =", out_txt)
199