[AutoTVM] [TVMC] Visualize tuning progress

[PhilippvK]

This change adds a callback for visualizing the FLOPS during tuning (AutoTVM & AutoScheduler (NEW)) over time using Matplotlib.

Usage

With Python

Create callback (multi=False: 1 window per Task, multi=True: all tasks in a single window):

visualize_callback = autotvm.callback.visualize_progress(task_idx, live=True, out_path="plot.png")

Pass callback to tuner:

tuner.tune(
  n_trial=100,
  measure_option=measure_option,
  callbacks=[visualize_callback],
)

With TVM command line

~~Just add a --visualize to the tvmc tune command.~~

There are three modes available:

• Live only: add --visualize live to tvmc tune command
• File only: add --visualize /pah/to/file.png to tvmc tune command
• Both: add --visualize live,/path/to/file.png to tvmc tune command

Screenshots

Additional changes

• Make si_prefix in tvmc/autotuner.py variable (Prerequisite for a later PR)

Open Questions

• ~~How to handle the matplotlib python dependency? Should I change python/gen_requirements.py and if yes which groups (tvmc/dev/?)?~~
• ~~The graph window automatically closes when the tuning is finished. Do we need an option to keep it open until closed?~~ Solved by newly introduced keep_open option
• ~~A similar feature for the AutoScheduler would be nice, too. I can look into it if there is interest.~~
• ~~Saving the plot to a file would also be possible. We only need to decide how to expose this to the cmdline (i.e. --visualize plot.png)~~

cc @Mousius @gromero

[PhilippvK]

CC @leandron @giuseros @Mousius @areusch @tqchen

[PhilippvK]

@leandron

@PhilippvK this looks very cool! the only thing missing is to declare the dependency on matplotlib so that the user is not surprised with and ImportError when trying to use it :).

I already had this in mind (see original post) but was unsure on what the best approach is.

I feel also this should be added to any of the tutorials, because it is too cool for users to miss out.

Yes. I will try to come up with a short section for the tutorials.

What do you think about the other questions?

• The graph window automatically closes when the tuning is finished. Do we need an option to keep it open until closed manually?
• A similar feature for the AutoScheduler would be nice, too. I can look into it if there is interest.
• Saving the plot to a file would also be possible. We only need to decide how to expose this to the cmdline (i.e. --visualize-out plot.png)

[PhilippvK]

I reworked the PR and updated the comment accordingly. Thank you for your inputs! Please have another look at the changes as I added several improvements:

• AutoScheduler support
• Mention in Tutorial
• Minimal unit test for cmdline parsing of --visualize argument
• Support file export

[PhilippvK]

I will followup with a new branch next year!