| # Licensed to the Apache Software Foundation (ASF) under one |
| # or more contributor license agreements. See the NOTICE file |
| # distributed with this work for additional information |
| # regarding copyright ownership. The ASF licenses this file |
| # to you under the Apache License, Version 2.0 (the |
| # "License"); you may not use this file except in compliance |
| # with the License. You may obtain a copy of the License at |
| # |
| # http://www.apache.org/licenses/LICENSE-2.0 |
| # |
| # Unless required by applicable law or agreed to in writing, |
| # software distributed under the License is distributed on an |
| # "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY |
| # KIND, either express or implied. See the License for the |
| # specific language governing permissions and limitations |
| # under the License. |
| |
| """Tuning a single group conv2d operator""" |
| |
| from collections import namedtuple |
| import logging |
| import os |
| |
| import tvm |
| from tvm import te |
| from tvm import autotvm |
| import topi |
| import vta |
| import vta.testing |
| |
| env = vta.get_env() |
| |
| Workload = namedtuple("GroupConv2DWorkload", |
| ['batch', 'height', 'width', 'in_filter', 'out_filter', 'groups', |
| 'hkernel', 'wkernel', 'hpad', 'wpad', 'hstride', 'wstride']) |
| |
| # Mobilenet (grouped variant) workloads |
| mobilenet_wkls = [ |
| ('mobilenet.D1', Workload(env.BATCH, 112, 112, 32, 32, 2, 3, 3, 1, 1, 1, 1)), |
| ('mobilenet.D2', Workload(env.BATCH, 112, 112, 64, 64, 4, 3, 3, 1, 1, 2, 2)), |
| ('mobilenet.D3', Workload(env.BATCH, 56, 56, 128, 128, 8, 3, 3, 1, 1, 1, 1)), |
| ('mobilenet.D4', Workload(env.BATCH, 56, 56, 128, 128, 8, 3, 3, 1, 1, 2, 2)), |
| ('mobilenet.D5', Workload(env.BATCH, 28, 28, 256, 256, 16, 3, 3, 1, 1, 1, 1)), |
| ('mobilenet.D6', Workload(env.BATCH, 28, 28, 256, 256, 16, 3, 3, 1, 1, 2, 2)), |
| ('mobilenet.D7', Workload(env.BATCH, 14, 14, 512, 512, 32, 3, 3, 1, 1, 1, 1)), |
| ('mobilenet.D8', Workload(env.BATCH, 14, 14, 512, 512, 32, 3, 3, 1, 1, 2, 2)), |
| ('mobilenet.D9', Workload(env.BATCH, 7, 7, 1024, 1024, 64, 3, 3, 1, 1, 1, 1)), |
| ] |
| |
| @tvm.te.tag_scope(tag=topi.tag.ELEMWISE) |
| def my_clip(x, a_min, a_max): |
| """Unlike topi's current clip, put min and max into two stages.""" |
| const_min = tvm.tir.const(a_min, x.dtype) |
| const_max = tvm.tir.const(a_max, x.dtype) |
| x = te.compute(x.shape, lambda *i: tvm.te.min(x(*i), const_max), name="clipA") |
| x = te.compute(x.shape, lambda *i: tvm.te.max(x(*i), const_min), name="clipB") |
| return x |
| |
| def group_conv2d(N, CI, H, W, CO, KH, KW, strides, padding, dilation, group): |
| |
| CI_G = CI // groups |
| data_shape = (N//env.BATCH, CI//env.BLOCK_IN, H, W, env.BATCH, env.BLOCK_IN) |
| kernel_shape = (CO//env.BLOCK_OUT, CI_G//env.BLOCK_IN, KH, KW, env.BLOCK_OUT, env.BLOCK_IN) |
| bias_shape = (N//env.BATCH, CO//env.BLOCK_OUT, 1, 1, env.BATCH, env.BLOCK_OUT) |
| |
| data = te.placeholder(data_shape, name="data", dtype=env.inp_dtype) |
| kernel = te.placeholder(kernel_shape, name="kernel", dtype=env.wgt_dtype) |
| bias = te.placeholder(bias_shape, name="bias", dtype=env.acc_dtype) |
| |
| with tvm.target.vta(): |
| res = topi.nn.group_conv2d_nchw( |
| data, |
| kernel, |
| strides, |
| padding, |
| dilation, |
| groups, |
| env.acc_dtype) |
| res = topi.right_shift(res, env.WGT_WIDTH) |
| res = topi.add(res, bias) |
| res = my_clip(res, 0, (1 << env.OUT_WIDTH - 1) - 1) |
| res = topi.cast(res, env.out_dtype) |
| |
| if tvm.target.Target.current().device_name == 'vta': |
| s = topi.generic.schedule_group_conv2d_nchw([res]) |
| else: |
| s = te.create_schedule([res.op]) |
| |
| return s, [data, kernel, bias, res] |
| |
| if __name__ == '__main__': |
| |
| # Logging config (for printing tuning log to the screen) |
| logging.basicConfig() |
| |
| # Tuning log files |
| log_file = "%s.group_conv2d.log" % (env.TARGET) |
| # create tmp log file |
| tmp_log_file = log_file + ".tmp" |
| if os.path.exists(log_file): |
| os.remove(log_file) |
| |
| # Get tracker info from env |
| tracker_host = os.environ.get("TVM_TRACKER_HOST", None) |
| tracker_port = os.environ.get("TVM_TRACKER_PORT", None) |
| if not tracker_host or not tracker_port: |
| print("Set your AutoTVM tracker node host and port variables to run the autotuner") |
| exit() |
| |
| for idx, (wl_name, wl) in enumerate(mobilenet_wkls): |
| prefix = "[Task %2d/%2d] " % (idx, len(mobilenet_wkls)) |
| |
| # Read in workload parameters |
| N = wl.batch |
| CI = wl.in_filter |
| H = wl.height |
| W = wl.width |
| CO = wl.out_filter |
| KH = wl.hkernel |
| KW = wl.wkernel |
| strides = (wl.hstride, wl.wstride) |
| padding = (wl.hpad, wl.wpad) |
| dilation = (1, 1) |
| groups = wl.groups |
| |
| # Create task |
| task = autotvm.task.create( |
| group_conv2d, |
| args=(N, CI, H, W, CO, KH, KW, strides, padding, dilation, groups), |
| target=tvm.target.vta(), |
| target_host=env.target_host, |
| template_key='direct') |
| print(task.config_space) |
| |
| # Tune |
| measure_option = autotvm.measure_option( |
| builder=autotvm.LocalBuilder(), |
| runner=autotvm.RPCRunner( |
| env.TARGET, host=tracker_host, port=int(tracker_port), |
| number=5, timeout=60, |
| check_correctness=True)) |
| |
| # Run Tuner |
| tuner = autotvm.tuner.RandomTuner(task) |
| tuner.tune( |
| n_trial=len(task.config_space), |
| early_stopping=None, |
| measure_option=measure_option, |
| callbacks=[ |
| autotvm.callback.progress_bar(len(task.config_space), prefix=prefix), |
| autotvm.callback.log_to_file(tmp_log_file)]) |
| |
| # Pick best records to a cache file |
| autotvm.record.pick_best(tmp_log_file, log_file) |
| os.remove(tmp_log_file) |