AutoTVM tree RNN model

dan · December 3, 2018, 2:18pm

I read the “Learning to Optimize Tensor Programs” paper, which describes the design of AutoTVM.
Though the focus is on the GBT cost model, I’m interested in replicating and exploring the TreeGRU model.

Does anyone have advice on embedding loop levels?
How I can access the tree structure of the AST to encode each level?
Can I leverage itervar features and curve sample features used by the GBT model?

merrymercy · December 11, 2018, 3:57am

Here is the related code. It is on my branch and not upstreamed.
You can find related tree-structure extraction code.

github.com

merrymercy/tvm/blob/autotvm-exp/python/tvm/autotvm/tuner/treegru_tuner.py

"""Tuner that uses xgboost as cost model"""

from .model_based_tuner import ModelBasedTuner, ModelOptimizer
from .treernn_cost_model import TreeRNNCostModel
from .sa_model_optimizer import SimulatedAnnealingOptimizer


class TreeGRUTuner(ModelBasedTuner):
    """Tuner that uses xgboost as cost model

    Parameters
    ----------
    task: Task
        The tuning task
    plan_size: int
        The size of a plan. After `plan_size` trials, the tuner will refit a new cost model
        and do planing for the next `plan_size` trials.
    feature_type: str, optional

    loss_type: str

This file has been truncated. show original

xwrock · January 8, 2020, 6:06am

Hi. I wonder why TreeGRU model was adopted in paper, but it is not on master branch of TVM? It seems that the TreeGRU works well according to the paper.