【9】TVM-5：用TVM编译Transformers模型、SSD模型、GCN模型-附python程序
这篇文档总结了TVM可以在cpu上编译部署的一些模型，包括 transformers模型、SSD模型、GCN模型 ，因为编译部署后还是应用fp32上进行推理，所以基本可以看做是文【4】【4-1】的系列文；
1. 利用TVM框架进行一个 Hugging Face 的修剪模型的cpu部署
整个过程如下：
1）下载一个 Transformers模型 ，并转化为tensorflow的GraphDef格式；
2）转换到TVM中的Relay图模型；
3）对上述模型编译，运行，推理一个随机输入，统计稠密图（Dense Graph）运行时间
4）对2）中的模型进行深度优化，剪枝，稀疏化，然后进行编译运行，推理同样一个随机输入，统计运行时间；
代码如下：
注：需要tensorflow版本2.2以上
# sphinx_gallery_start_ignore
from tvm import testing
testing.utils.install_request_hook(depth=3)
# sphinx_gallery_end_ignore
###############################################################################
# Load Required Modules
# ---------------------
# Other than TVM, scipy, the latest transformers, and
# tensorflow 2.2+ are required.
import os
import tvm
import time
import itertools
import numpy as np
import tensorflow as tf
from tvm import relay, runtime
from tvm.contrib import graph_executor
from tvm.relay import data_dep_optimization as ddo
from tensorflow.python.framework.convert_to_constants import (
convert_variables_to_constants_v2,
import scipy.sparse as sp
# Ask tensorflow to limit its GPU memory to what's actually needed
# instead of gobbling everything that's available.
# https://www.tensorflow.org/guide/gpu#limiting_gpu_memory_growth
# This way this tutorial is a little more friendly to sphinx-gallery.
gpus = tf.config.list_physical_devices("GPU")
if gpus:
try:
for gpu in gpus:
tf.config.experimental.set_memory_growth(gpu, True)
print("tensorflow will use experimental.set_memory_growth(True)")
except RuntimeError as e:
print("experimental.set_memory_growth option is not available: {}".format(e))
###############################################################################
# Configure Settings
# ------------------
# Let's start by defining some parameters that define the type of model
# and sparsity to run.
# The name of the transformer model to download and run.
name = "huggingface/prunebert-base-uncased-6-finepruned-w-distil-squad"
# The number of batches in an input.
batch_size = 1
# The length of each input sequence.
seq_len = 128
# TVM platform identifier. Note that best cpu performance can be achieved by setting -mcpu
# appropriately for your specific machine. CUDA and ROCm are also supported.
target = "llvm"
# Which device to run on. Should be one of tvm.cpu() or tvm.cuda().
dev = tvm.cpu()
# If true, then a sparse variant of the network will be run and
# benchmarked.
measure_sparse = True
# The block size of structured sparsity to convert weight tensors
# into. Changing this parameter may yield speedups for some platforms.
bs_r = 1
# For models besides PruneBert (which is 95% sparse), this parameter
# determines how sparse the generated weights should be. The higher
# the sparsity, the faster the result.
sparsity = 0.85
###############################################################################
# Download and Convert Transformers Model
# ---------------------------------------
# Now we'll grab a model from the transformers module, download it,
# convert it into a TensorFlow graphdef in preperation for converting that graphdef into
# a relay graph that we can optimize and deploy.
def load_keras_model(module, name, seq_len, batch_size, report_runtime=True):
model = module.from_pretrained(name)
dummy_input = tf.keras.Input(shape=[seq_len], batch_size=batch_size, dtype="int32")
dummy_out = model(dummy_input) # Propagate shapes through the keras model.
if report_runtime:
np_input = np.random.uniform(size=[batch_size, seq_len], low=0, high=seq_len).astype("int32")
start = time.time()
repeats = 50
for i in range(repeats):
np_out = model(np_input)
end = time.time()
print("Keras Runtime: %f ms." % (1000 * ((end - start) / repeats)))
return model
def convert_to_graphdef(model, batch_size, seq_len):
model_func = tf.function(lambda x: model(x))
input_dict = model._saved_model_inputs_spec
input_spec = input_dict[list(input_dict.keys())[0]]
model_func = model_func.get_concrete_function(
tf.TensorSpec([batch_size, seq_len], input_spec.dtype)
frozen_func = convert_variables_to_constants_v2(model_func)
return frozen_func.graph.as_graph_def()
def download_model(name, batch_size, seq_len):
import transformers
module = getattr(transformers, "TFBertForSequenceClassification")
model = load_keras_model(module, name=name, batch_size=batch_size, seq_len=seq_len)
return convert_to_graphdef(model, batch_size, seq_len)
###############################################################################
# Convert to Relay Graph
# ----------------------
# We now have all the tooling to get a transformers model in the right format
# for relay conversion. Let's import it! In the following function we
# save the imported graph in relay's json format so that we dont have
# to reimport from tensorflow each time this script is run.
def import_graphdef(
name,
batch_size,
seq_len,
save_relay=True,
relay_file="model.json",
relay_params="model.params",
abs_path = os.path.dirname(os.path.abspath(__file__))
shape_dict = {"input_1": (batch_size, seq_len)}
relay_file = ("%s_%d_%d_%s" % (name, batch_size, seq_len, relay_file)).replace("/", "_")
relay_params = ("%s_%d_%d_%s" % (name, batch_size, seq_len, relay_params)).replace("/", "_")
if os.path.exists(os.path.join(abs_path, relay_file)) and os.path.exists(
os.path.join(abs_path, relay_params)
with open(os.path.join(abs_path, relay_file), "r") as fi:
mod = tvm.ir.load_json(fi.read())
with open(os.path.join(abs_path, relay_params), "rb") as fi:
params = relay.load_param_dict(fi.read())
else:
graph_def = download_model(name, batch_size, seq_len)
mod, params = relay.frontend.from_tensorflow(graph_def, shape=shape_dict)
if save_relay:
with open(os.path.join(abs_path, relay_file), "w") as fo:
fo.write(tvm.ir.save_json(mod))
with open(os.path.join(abs_path, relay_params), "wb") as fo:
fo.write(runtime.save_param_dict(params))
return mod, dict(params.items()), shape_dict
###############################################################################
# Run the Dense Graph
# -------------------
# Let's run the default version of the imported model. Note that even if
# the weights are sparse, we won't see any speedup because we are using
# regular dense matrix multiplications on these dense (but mostly zero)
# tensors instead of sparse aware kernels.
def
run_relay_graph(mod, params, shape_dict, target, dev):
with relay.build_config(opt_level=3):
lib = relay.build(mod, target=target, params=params)
input_shape = shape_dict["input_1"]
dummy_data = np.random.uniform(size=input_shape, low=0, high=input_shape[1]).astype("int32")
m = graph_executor.GraphModule(lib["default"](dev))
m.set_input(0, dummy_data)
m.run()
tvm_output = m.get_output(0)
print(m.benchmark(dev, repeat=5, number=5))
return tvm_output
def run_dense(mod, params, shape_dict, target, dev):
print("Dense Model Benchmark:")
return run_relay_graph(mod, params, shape_dict, target, dev)
###############################################################################
# Run the Sparse Graph
# --------------------
# Next we'll convert the graph into a sparse representation and generate
# fake sparse weights if needed. Then we'll use the same benchmarking
# script as dense to see how much faster we go! We apply a few relay passes
# to the graph to get it leveraging sparsity. First we use
# `simplify_fc_transpose` to use transposes on the weights of dense layers
# into the parameters. This makes it easier to convert to matrix multiplies
# to sparse versions. Next we apply `bsr_dense.convert` to identify all
# weight matrices that can be sparse, and automatically replace them.
# The `bsr_dense.convert` call below is doing the heavy lifting of identifying
# which weights in the model can be made sparse by checking if they are
# at least `sparsity_threshold` percent sparse. If so, it converts those
# weights into *Block Compressed Row Format (BSR)*. BSR is essentially
# a representation that indexes into the nonzero chunks of the tensor,
# making it easy for an algorithm to load those non-zero chunks and ignore
# the rest of the tensor. Once the sparse weights are in BSR format,
# `relay.transform.DenseToSparse` is applied to actually replace
# `relay.dense` operations with `relay.sparse_dense` calls that can be
# run faster.
def random_bsr_matrix(M, N, BS_R, BS_C, density, dtype="float32"):
Y = np.zeros((M, N), dtype=dtype)
assert M % BS_R == 0
assert N % BS_C == 0
nnz = int(density * M * N)
num_blocks = int(nnz / (BS_R * BS_C)) + 1
candidate_blocks = np.asarray(list(itertools.product(range(0, M, BS_R), range(0, N, BS_C))))
assert candidate_blocks.shape[0] == M // BS_R * N // BS_C
chosen_blocks = candidate_blocks[
np.random.choice(candidate_blocks.shape[0], size=num_blocks, replace=False)
for i in range(len(chosen_blocks)):
r, c = chosen_blocks[i]
Y[r : r + BS_R, c : c + BS_C] = np.random.uniform(-0.1, 0.1, (BS_R, BS_C))
s = sp.bsr_matrix(Y, blocksize=(BS_R, BS_C))
assert s.data.shape == (num_blocks, BS_R, BS_C)
assert s.data.size >= nnz
assert s.indices.shape == (num_blocks,)
assert s.indptr.shape == (M // BS_R + 1,)
return s.todense()
def random_sparse_bert_params(func, params, density, BS_R, BS_C):
def deepcopy(param_dic):
ret = {}
for k, v in param_dic.items():
ret[k] = tvm.nd.array(v.numpy())
return ret
new_params = deepcopy(params)
dense_weight_names = relay.analysis.sparse_dense._search_dense_op_weight(func)
for item in dense_weight_names:
name = str(item)
shape = new_params[name].shape
if shape[0] % BS_R == 0 and shape[1] % BS_C == 0:
new_w = random_bsr_matrix(shape[0], shape[1], BS_R, BS_C, density)
new_params[name] = tvm.nd.array(new_w)
return new_params
def run_sparse(mod, params, shape_dict, target, dev, bs_r, sparsity, gen_weights):
mod, params = ddo.simplify_fc_transpose.convert(mod["main"], params)
if gen_weights:
params = random_sparse_bert_params(mod, params, BS_R=bs_r, BS_C=1, density=1 - sparsity)
mod, params = ddo.bsr_dense.convert(mod, params, (bs_r, 1), sparsity_threshold=0.8)
print("Block Sparse Model with {blocksize}x1 blocks:".format(
blocksize=bs_r))
return run_relay_graph(mod, params, shape_dict, target, dev)
###############################################################################
# Run All the Code!
# -----------------
# And that's it! Now we'll simply call all the needed function to benchmark
# the model according to the set parameters. Note that to run this code
# you'll need to uncomment the last line first.
def benchmark():
mod, params, shape_dict = import_graphdef(name, batch_size, seq_len)
run_dense(mod, params, shape_dict, target, dev)
if measure_sparse:
gen_weights = "prune" not in name
run_sparse(mod, params, shape_dict, target, dev, bs_r, sparsity, gen_weights)
if __name__ == "__main__":
benchmark()
结果输出稠密模型的推理时间与稀疏模型的推理时间比较：
结论：经过TVM稀疏化的Transformers模型部署后，模型推理速度相比优化前速度提升近一倍；
完整代码见： https:// github.com/Learneruestc /TVM-all/blob/main/TVM4/deploy_sparse.py
2. 利用TVM框架进行一个 SSD模型的cpu部署
整个过程如下：
1）下载一个GluonCV预训练好的SSD模型，并转化为Relay图模型；
2）对上述模型编译，运行，推理一个例子图，输出检测目标的标记框及置信度
代码如下：
# sphinx_gallery_start_ignore
from tvm import testing
testing.utils.install_request_hook(depth=3)
# sphinx_gallery_end_ignore
import tvm
from tvm import te
from matplotlib import pyplot as plt
from tvm import relay
from tvm.contrib import graph_executor
from tvm.contrib.download import download_testdata
from gluoncv import model_zoo, data, utils
######################################################################
# Preliminary and Set parameters
# ------------------------------
# .. note::
# We support compiling SSD on both CPUs and GPUs now.
# To get best inference performance on CPU, change
# target argument according to your device and
# follow the :ref:`tune_relay_x86` to tune x86 CPU and
# :ref:`tune_relay_arm` for arm CPU.
# To get best inference performance on Intel graphics,
# change target argument to :code:`opencl -device=intel_graphics`.
# But when using Intel graphics on Mac, target needs to
# be set to `opencl` only for the reason that Intel subgroup
# extension is not supported on Mac.
# To get best inference performance on CUDA-based GPUs,
# change the target argument to :code:`cuda`; and for
# OPENCL-based GPUs, change target argument to
# :code:`opencl` followed by device argument according
# to your device.
supported_model = [
"ssd_512_resnet50_v1_voc",
"ssd_512_resnet50_v1_coco",
"ssd_512_resnet101_v2_voc",
"ssd_512_mobilenet1.0_voc",
"ssd_512_mobilenet1.0_coco",
"ssd_300_vgg16_atrous_voc" "ssd_512_vgg16_atrous_coco",
model_name = supported_model[0]
dshape = (1, 3, 512, 512)
######################################################################
# Download and pre-process demo image
im_fname = download_testdata(
"https://github.com/dmlc/web-data/blob/main/" + "gluoncv/detection/street_small.jpg?raw=true",
"street_small.jpg",
module="data",
x, img = data.transforms.presets.ssd.load_test(im_fname, short=512)
######################################################################
# Convert and compile model for CPU.
block = model_zoo.get_model(model_name, pretrained=True)
def build(target):
mod, params = relay.frontend.from_mxnet(block, {"data": dshape})
with tvm.transform.PassContext(opt_level=3):
lib = relay.build(mod, target, params=params)
return lib
######################################################################
# Create TVM runtime and do inference
# .. note::
# Use target = "cuda -libs" to enable thrust based sort, if you
# enabled thrust during cmake by -DUSE_THRUST=ON.
def run(lib, dev):
# Build TVM runtime
m = graph_executor.GraphModule(lib["default"](dev))
tvm_input = tvm.nd.array(x.asnumpy(), device=dev)
m.set_input("data", tvm_input)
# execute
m.run()
# get outputs
class_IDs, scores, bounding_boxs = m.get_output(0), m.get_output(1), m.get_output(2)
return class_IDs, scores, bounding_boxs
for target in ["llvm", "cuda"]:
dev = tvm.device(target, 0)
if dev.exist:
lib = build(target)
class_IDs, scores, bounding_boxs = run(lib, dev)
######################################################################
# Display result
ax = utils.viz.plot_bbox(
bounding_boxs.numpy()[0],
scores.numpy()[0],
class_IDs.numpy()[0],
class_names=block.classes,
plt.show()
输出结果是tvm编译后的模型能够准备检测出图片中目标并打框、显示置信度：
完整代码见： https:// github.com/Learneruestc /TVM-all/blob/main/TVM4/deploy_ssd_gluoncv.py
3. 利用TVM框架进行图卷积神经网络GCN的部署
整个过程如下：
1）下载一个预训练好的 DGL-PyTorch模型，用Cora数据集进行测试，输出前5个输出和测试准确率；
2）在TVM的Relay里构建图神经网络，对模型进行编译、推理，输出前5个输出和测试准确率；
代码如下：
注：需安装dgl版本为0.6.1, pip install dgl==0.6.1
######################################################################
# Define GCN in DGL with PyTorch backend
# --------------------------------------
# DGL example: https://github.com/dmlc/dgl/tree/master/examples/pytorch/gcn
# This part reuses the code from the above example.
import torch
import torch.nn as nn
import torch.nn.functional as F
import dgl
import networkx as nx
from dgl.nn.pytorch import GraphConv
class GCN(nn.Module):
def __init__(self, g, n_infeat, n_hidden, n_classes, n_layers, activation):
super(GCN, self).__init__()
self.g = g
self.layers = nn.ModuleList()
self.layers.append(GraphConv(n_infeat, n_hidden, activation=activation))
for i in range(n_layers - 1):
self.layers.append(GraphConv(n_hidden, n_hidden, activation=activation))
self.layers.append(GraphConv(n_hidden, n_classes))
def forward(self, features):
h = features
for i, layer in enumerate(self.layers):
# handle api changes for differnt DGL version
if dgl.__version__ > "0.3":
h = layer(self.g, h)
else:
h = layer(h, self.g)
return h
######################################################################
# Define the functions to load dataset and evaluate accuracy
# ----------------------------------------------------------
# You may substitute this part with your own dataset, here we load data from DGL
from dgl.data import load_data
from collections import namedtuple
def load_dataset(dataset="cora"):
args = namedtuple("args", ["dataset"])
data = load_data(args(dataset))
# Remove self-loops to avoid duplicate passing of a node's feature to itself
g = data.graph
g.remove_edges_from(nx.selfloop_edges(g))
g.add_edges_from(zip(g.nodes, g.nodes))
return g, data
def evaluate(data, logits):
test_mask = data.test_mask # the test set which isn't included in the training phase
pred = logits.argmax(axis=1)
acc = ((pred == data.labels) * test_mask).sum() / test_mask.sum()
return acc
######################################################################
# Load the data and set up model parameters
# -----------------------------------------
Parameters
----------
dataset: str
Name of dataset. You can choose from ['cora', 'citeseer', 'pubmed'].
num_layer: int
number of hidden layers
num_hidden: int
number of the hidden units in the hidden layer
infeat_dim: int
dimension of the input features
num_classes: int
dimension of model output (Number of classes)
# sphinx_gallery_start_ignore
from tvm import testing
testing.utils.install_request_hook(depth=3)
# sphinx_gallery_end_ignore
dataset = "cora"
g, data = load_dataset(dataset)
num_layers = 1
num_hidden = 16
infeat_dim = data.features.shape[1]
num_classes = data.num_labels
######################################################################
# Set up the DGL-PyTorch model and get the golden results
# -------------------------------------------------------
# The weights are trained with https://github.com/dmlc/dgl/blob/master/examples/pytorch/gcn/train.py
from tvm.contrib.download import download_testdata
from dgl import DGLGraph
features = torch.FloatTensor(data.features)
dgl_g = DGLGraph(g)
torch_model = GCN(dgl_g, infeat_dim, num_hidden, num_classes, num_layers, F.relu)
# Download the pretrained weights
model_url = "https://homes.cs.washington.edu/~cyulin/media/gnn_model/gcn_%s.torch" % (dataset)
model_path = download_testdata(model_url, "gcn_%s.pickle" % (dataset), module="gcn_model")
# Load the weights into the model
torch_model.load_state_dict(torch.load(model_path))
######################################################################
# Run the DGL model and test for accuracy
# ---------------------------------------
torch_model.eval()
with torch.no_grad():
logits_torch = torch_model(features)
print("Print the first five outputs from DGL-PyTorch execution\n", logits_torch[:5])
acc = evaluate(data, logits_torch.numpy())
print("Test accuracy of DGL results: {:.2%}".format(acc))
######################################################################
# Define Graph Convolution Layer in Relay
# ---------------------------------------
# To run GCN on TVM, we first need to implement Graph Convolution Layer.
# You may refer to https://github.com/dmlc/dgl/blob/master/python/dgl/nn/mxnet/conv/graphconv.py for a GraphConv Layer implemented in DGL with MXNet Backend
# The layer is defined with below operations, note that we apply two transposes to keep adjacency matrix on right hand side of sparse_dense operator,
# this method is temporary and will be updated in next few weeks when we have sparse matrix transpose and support for left sparse operator.
# .. math::
# \mbox{GraphConv}(A, H, W) = A * H * W
# = ((H * W)^t * A^t)^t
# = ((W^t * H^t) * A^t)^t
from tvm import relay
from tvm.contrib import graph_executor
import tvm
from tvm import te
def GraphConv(layer_name, input_dim, output_dim, adj, input, norm=None, bias=True, activation=None):
Parameters
----------
layer_name: str
Name of layer
input_dim: int
Input dimension per node feature
output_dim: int,
Output dimension per node feature
adj: namedtuple,
Graph representation (Adjacency Matrix) in Sparse Format (`data`, `indices`, `indptr`),
where `data` has shape [num_nonzeros], indices` has shape [num_nonzeros], `indptr` has shape [num_nodes + 1]
input: relay.Expr,
Input feature to current layer with shape [num_nodes, input_dim]
norm: relay.Expr,
Norm passed to this layer to normalize features before and after Convolution.
bias: bool
Set bias to True to add bias when doing GCN layer
activation: <function relay.op.nn>,
Activation function applies to the output. e.g. relay.nn.{relu, sigmoid, log_softmax, softmax, leaky_relu}
Returns
----------
output: tvm.relay.Expr
The Output Tensor for this layer [num_nodes, output_dim]
if norm is not None:
input = relay.multiply(input, norm)
weight = relay.var(layer_name + ".weight", shape=(input_dim, output_dim))
weight_t = relay.transpose(weight)
dense = relay.nn.dense(weight_t, input)
output = relay.nn.sparse_dense(dense, adj)
output_t = relay.transpose(output)
if norm is not None:
output_t = relay.multiply(output_t, norm)
if bias is True:
_bias = relay.var(layer_name + ".bias", shape=(output_dim, 1))
output_t = relay.nn.bias_add(output_t, _bias, axis=-1)
if activation is not None:
output_t = activation(output_t)
return output_t
######################################################################
# Prepare the parameters needed in the GraphConv layers
# -----------------------------------------------------
import numpy as np
import networkx as nx
def prepare_params(g, data):
params = {}
params["infeats"] = data.features.numpy().astype(
"float32"
) # Only support float32 as feature for now
# Generate adjacency matrix
adjacency = nx.to_scipy_sparse_matrix(g)
params["g_data"] = adjacency.data.astype("float32")
params["indices"] = adjacency.indices.astype("int32")
params["indptr"] = adjacency.indptr.astype("int32")
# Normalization w.r.t. node degrees
degs = [g.in_degree[i] for i in range(g.number_of_nodes())]
params["norm"] = np.power(degs, -0.5).astype("float32")
params["norm"] = params["norm"].reshape((params["norm"].shape[0], 1))
return params
params = prepare_params(g, data)
# Check shape of features and the validity of adjacency matrix
assert len(params["infeats"].shape) == 2
assert (
params["g_data"] is not None and params["indices"] is not None and params["indptr"] is not None
assert params["infeats"].shape[0] == params["indptr"].shape[0] - 1
######################################################################
# Put layers together
# -------------------
# Define input features, norms, adjacency matrix in Relay
infeats = relay.var("infeats", shape=data.features.shape)
norm = relay.Constant(tvm.nd.array(params["norm"]))
g_data = relay.Constant(tvm.nd.array(params["g_data"]))
indices = relay.Constant(tvm.nd.array(params["indices"]))
indptr = relay.Constant(tvm.nd.array(params["indptr"]))
Adjacency = namedtuple("Adjacency", ["data", "indices", "indptr"])
adj = Adjacency(g_data, indices, indptr)
# Construct the 2-layer GCN
layers = []
layers.append(
GraphConv(
layer_name="layers.0",
input_dim=infeat_dim,
output_dim=num_hidden,
adj=adj,
input=infeats,
norm=norm,
activation=relay.nn.relu,
layers.append(
GraphConv(
layer_name="layers.1",
input_dim=num_hidden,
output_dim=num_classes,
adj=adj,
input=layers[-1],
norm=norm,
activation=None,
# Analyze free variables and generate Relay function
output = layers[-1]
######################################################################
# Compile and run with TVM
# ------------------------
# Export the weights from PyTorch model to Python Dict
model_params = {}
for param_tensor in torch_model.state_dict():
model_params[param_tensor] = torch_model.state_dict()[param_tensor].numpy()
for i in range(num_layers + 1):
params["layers.%d.weight" % (i)] = model_params["layers.%d.weight" % (i)]
params["layers.%d.bias" % (i)] = model_params["layers.%d.bias" % (i)]
# Set the TVM build target
target = "llvm" # Currently only support `llvm` as target
func = relay.Function(relay.analysis.free_vars(output), output)
func = relay.build_module.bind_params_by_name(func, params)
mod = tvm.IRModule()
mod["main"] = func
# Build with Relay
with tvm.transform.PassContext(opt_level=0): # Currently only support opt_level=0
lib = relay.build(mod, target, params=params)
# Generate graph executor
dev = tvm.device(target, 0)
m = graph_executor.GraphModule(lib["default"](dev))
######################################################################
# Run the TVM model, test for accuracy and verify with DGL
# --------------------------------------------------------
m.run()
logits_tvm = m.get_output(0).numpy()
print("Print the first five outputs from TVM execution\n", logits_tvm[:5])
labels = data.labels
test_mask = data.test_mask
acc = evaluate(data, logits_tvm)