Advertisement
Not a member of Pastebin yet?
Sign Up,
it unlocks many cool features!
- {
- "cells": [
- {
- "cell_type": "code",
- "execution_count": 1,
- "metadata": {},
- "outputs": [],
- "source": [
- "import numpy as np\n",
- "\n",
- "from tvm import relay\n",
- "from tvm.relay import testing\n",
- "from tvm.relay.testing import layers\n",
- "from tvm.relay.op.nn import nn\n",
- "\n",
- "from tvm.relay.testing.init import create_workload\n",
- "import tvm\n",
- "from tvm.contrib import graph_runtime"
- ]
- },
- {
- "cell_type": "code",
- "execution_count": 2,
- "metadata": {},
- "outputs": [
- {
- "name": "stdout",
- "output_type": "stream",
- "text": [
- "26 26\n"
- ]
- }
- ],
- "source": [
- "# set hyperparams\n",
- "batch_size = 1\n",
- "in_c = 512\n",
- "num_filters = 512\n",
- "in_h, in_w = 28, 28\n",
- "kdim = 3\n",
- "groups=1\n",
- "stride=1\n",
- "pad=0\n",
- "data_shape = (batch_size, in_c, in_h, in_w)\n",
- "num_groups = [1, 2, 4, 8, 16, 32, 64, 128, 256, 512]\n",
- "out_h = (in_h + 2* pad - kdim) // stride + 1\n",
- "out_w = (in_w + 2* pad - kdim) // stride + 1\n",
- "print(out_w, out_h)\n",
- "out_shape = (1, num_filters, out_h, out_w)"
- ]
- },
- {
- "cell_type": "code",
- "execution_count": 3,
- "metadata": {},
- "outputs": [],
- "source": [
- "def conv_block(data, name, num_filters, kernel_size=(3, 3), strides=(1, 1),\n",
- " padding=(1, 1), groups=1, dtype='float32'):\n",
- " \"\"\"Helper function to construct conv layer\"\"\"\n",
- " kdim1, kdim2 = kernel_size\n",
- " weight = relay.var(name + \"_weight\")\n",
- " conv = nn.conv2d(\n",
- " data=data,\n",
- " weight=weight,\n",
- " channels=num_filters,\n",
- " kernel_size=kernel_size,\n",
- " strides=strides,\n",
- " padding=padding,\n",
- " groups=groups)\n",
- " return conv"
- ]
- },
- {
- "cell_type": "code",
- "execution_count": 4,
- "metadata": {},
- "outputs": [],
- "source": [
- "def weenet(data_shape, num_filters, kdim, stride,\n",
- " pad, groups, dtype='float32'):\n",
- " \"\"\"Function to construct a WeeNet\"\"\"\n",
- " data = relay.var(\"data\", shape=data_shape, dtype=dtype)\n",
- " body = conv_block(data, 'conv_block_1', num_filters, (kdim,kdim), \n",
- " (stride,stride), (pad,pad), groups, dtype)\n",
- " return relay.Function(relay.ir_pass.free_vars(body), body)"
- ]
- },
- {
- "cell_type": "code",
- "execution_count": 5,
- "metadata": {},
- "outputs": [],
- "source": [
- "def get_workload(data_shape, num_filters, kdim, stride,\n",
- " pad, groups, dtype='float32'):\n",
- " net = weenet(data_shape, num_filters, kdim, stride, pad, groups, dtype)\n",
- " return create_workload(net)"
- ]
- },
- {
- "cell_type": "code",
- "execution_count": 6,
- "metadata": {},
- "outputs": [],
- "source": [
- "def compile_model(net, params, data, opt_level=1, ctx=tvm.cpu(0), target='llvm'):\n",
- " with relay.build_config(opt_level=opt_level):\n",
- " graph, lib, params = relay.build_module.build(\n",
- " net, target, params=params)\n",
- " # create module\n",
- " module = graph_runtime.create(graph, lib, ctx)\n",
- " # set input and parameters\n",
- " module.set_input(\"data\", data)\n",
- " module.set_input(**params)\n",
- " # run\n",
- " module.run()\n",
- " # get output\n",
- " out = module.get_output(0, tvm.nd.empty(out_shape)).asnumpy()\n",
- " return out, module"
- ]
- },
- {
- "cell_type": "code",
- "execution_count": 7,
- "metadata": {},
- "outputs": [],
- "source": [
- "# create random input\n",
- "np.random.seed(0)\n",
- "data_shape = (batch_size, in_c, in_h, in_w)\n",
- "data = np.random.uniform(-1, 1, size=data_shape).astype(\"float32\")\n"
- ]
- },
- {
- "cell_type": "code",
- "execution_count": 8,
- "metadata": {},
- "outputs": [
- {
- "name": "stderr",
- "output_type": "stream",
- "text": [
- "Cannot find config for target=llvm, workload=('conv2d', (1, 512, 28, 28, 'float32'), (512, 512, 3, 3, 'float32'), (1, 1), (0, 0), (1, 1), 'NCHW', 'float32'). A fallback configuration is used, which may bring great performance regression.\n",
- "WARNING:autotvm:Cannot find config for target=llvm, workload=('depthwise_conv2d_nchw', (1, 512, 28, 28, 'float32'), (512, 1, 3, 3, 'float32'), (1, 1), (0, 0), (1, 1), 'float32'). A fallback configuration is used, which may bring great performance regression.\n"
- ]
- }
- ],
- "source": [
- "results = [None] * len(num_groups)\n",
- "ctx=tvm.cpu(0)\n",
- "opt_level = 1\n",
- "\n",
- "for i, g in enumerate(num_groups):\n",
- " net, params = get_workload(data_shape, num_filters, kdim, stride, pad, g)\n",
- " out, m = compile_model(net, params, data, opt_level, ctx)\n",
- " f = m.module.time_evaluator('run', ctx)\n",
- " results[i] = f().mean\n",
- " "
- ]
- },
- {
- "cell_type": "code",
- "execution_count": 9,
- "metadata": {},
- "outputs": [
- {
- "data": {
- "text/plain": [
- "[0.22498705270000002,\n",
- " 0.8115864547999999,\n",
- " 0.4177320249,\n",
- " 0.2104043008,\n",
- " 0.1027479518,\n",
- " 0.0514274909,\n",
- " 0.0266668331,\n",
- " 0.0042375367,\n",
- " 0.0018110453,\n",
- " 0.0008504548000000001]"
- ]
- },
- "execution_count": 9,
- "metadata": {},
- "output_type": "execute_result"
- }
- ],
- "source": [
- "results"
- ]
- },
- {
- "cell_type": "code",
- "execution_count": null,
- "metadata": {},
- "outputs": [],
- "source": []
- }
- ],
- "metadata": {
- "kernelspec": {
- "display_name": "meth",
- "language": "python",
- "name": "meth"
- },
- "language_info": {
- "codemirror_mode": {
- "name": "ipython",
- "version": 3
- },
- "file_extension": ".py",
- "mimetype": "text/x-python",
- "name": "python",
- "nbconvert_exporter": "python",
- "pygments_lexer": "ipython3",
- "version": "3.5.3"
- }
- },
- "nbformat": 4,
- "nbformat_minor": 2
- }
Advertisement
Add Comment
Please, Sign In to add comment
Advertisement