How to filter conv2d workloads when using ansor tune

1

Hi friends: I have a question. Suppose I have a model which has many conv2d and other kinds of ops, may I ask whether I can only extract conv2d-related tasks (ansor-tune)?

Thanks!

2

You can use search for "conv2d" in the str(task.compute_dag).

3

If your TVM is new enough (after April), you can also use the following logic:

tasks, weights = extract_tasks(...)
for task, weight in zip(tasks, weights):
  if "conv2d" not in task.desc:
    # throw away this task
1 Like
4

great! thank you, I will try!