[TVM PyTorch Integration] optimized_torch & as_torch how-to guide #12318

juda · 2022-08-05T12:33:51Z

This PR provides two how-to guides to show the usage of

optimized_torch: tuning a PyTorch model/function by MetaSchedule
as_torch: wrap TVMscript into a PyTorch model/function

@yelite @junrushao1994 @masahi

gallery/how_to/work_with_pytorch/using_as_torch.py

masahi · 2022-08-17T08:35:52Z

gallery/how_to/work_with_pytorch/using_as_torch.py

+# Write your own PyTorch operator by TVMscript
+# -------------------------------
+# PyTorch is a very popular machine learning framework in which
+# it highly optimizes most commonly used operators.


PyTorch is a very popular machine learning framework which contains optimized implementations of most commonly used operators

masahi · 2022-08-17T08:43:06Z

gallery/how_to/work_with_pytorch/using_as_torch.py

+# PyTorch is a very popular machine learning framework in which
+# it highly optimizes most commonly used operators.
+# Nevertheless, sometimes you might want to write your own operators
+# in PyTorch, but the performance could be not satisfactory.


Nevertheless, sometimes you might want to write your own operators in PyTorch. In that case, the performance of such custom operators might not be satisfactory for your needs.

masahi · 2022-08-17T08:43:39Z

gallery/how_to/work_with_pytorch/using_as_torch.py

+# For example, assume you are writing a variance of MobileNet,
+# and you need to define a 1-d depthwise convolution operator.
+# Assume the number of in_channel and out_channel are both 700,
+# the width is 800 and the kernel size is 50,


the code uses kernel size 20

masahi · 2022-08-17T08:44:15Z

gallery/how_to/work_with_pytorch/using_as_torch.py

+
+def torch_depthwise(inputs, filters):
+    global out_channel
+    global kernel_size


do you need global here?

masahi · 2022-08-17T08:45:16Z

gallery/how_to/work_with_pytorch/using_as_torch.py

+# Nevertheless, sometimes you might want to write your own operators
+# in PyTorch, but the performance could be not satisfactory.
+#
+# For example, assume you are writing a variance of MobileNet,


It doesn't make a lot of sense to talk about a variant of Mobilenet (where only 2d convolution is used) but then suddenly bring up 1D convolution.

masahi · 2022-08-17T08:56:26Z

gallery/how_to/work_with_pytorch/using_as_torch.py

+)
+
+# We can tune the TVMscript code by providing a target device.
+# The model will deploy on CPU, and the optimization (e.g. tiling) will conduct automatically.


There are basic grammar issues here.

Instead of "deploy", just use "run". But more importantly, we want to say that the model will be "tuned" for CPU.

masahi · 2022-08-17T08:59:48Z

gallery/how_to/work_with_pytorch/using_as_torch.py

+
+print(tvm_depthwise.script())
+
+# Hint: If user plan to deploy on GPU, the GPU target should be provided,


Again, there is no "target" in this tutorial.

masahi · 2022-08-17T09:00:48Z

gallery/how_to/work_with_pytorch/using_as_torch.py

+print(tvm_depthwise.script())
+
+# Hint: If user plan to deploy on GPU, the GPU target should be provided,
+# and all the PyTorch tensors should convert into GPU version.


This sentence alone doesn't make sense.

masahi · 2022-08-17T09:02:28Z

gallery/how_to/work_with_pytorch/using_as_torch.py

+
+# In the working machine, the average inference time of `tvm_depthwise` is 120.0 us (TVM version is 0.9.0),
+# while the average inference time of `torch_depthwise` is 210.0 us (PyTorch version is 1.11.0),
+# showing the performance arises by around 43%.


showing the speedup of around 43%

masahi · 2022-08-17T09:03:08Z

gallery/how_to/work_with_pytorch/using_as_torch.py

+compare = benchmark.Compare(results)
+compare.print()
+
+# In the working machine, the average inference time of `tvm_depthwise` is 120.0 us (TVM version is 0.9.0),


In author's environment,

yelite · 2022-08-29T04:44:01Z

gallery/how_to/work_with_pytorch/using_as_torch.py

+@as_torch
+@T.prim_func
+def tvm_depthwise(
+    A: T.Buffer((70, 80), "float32"),


Will it be useful to show a follow-up example of how to make shape a variable by having nested function?

If we pass shape variables then we need match_buffer operators, which might confuse users.
Currently, I choose a minimal set of grammar.
@masahi what's your idea?

I think the Buffer syntax sugar can be extended for dynamic shapes. But currently we cannot tune over dynamic shapes, so the performance will probably be slower than PT.

yelite · 2022-08-29T04:47:15Z

gallery/how_to/work_with_pytorch/using_as_torch.py

+# We can build the TVMscript code by calling the `tune` method.
+# Without providing more information, the model will be tuned for CPU.
+
+tvm_depthwise.tune()


Will it be better to explicitly write down the default TuneConfig and target here so that reader has better idea on how to customize this?

juda · 2022-09-01T01:14:49Z

@masahi I have improved the text, could you please review it again?

masahi · 2022-09-01T01:50:32Z

gallery/how_to/work_with_pytorch/using_as_torch.py

+# Nevertheless, sometimes you might want to write your own operators in PyTorch.
+# In that case, the performance of such custom operators might not be satisfactory for your needs.
+#
+# One of the examples is to define a 1-d depthwise convolution operator.


"For example, suppose we want to define..."

masahi · 2022-09-01T01:55:28Z

gallery/how_to/work_with_pytorch/using_as_torch.py

+
+# Then, we plan to optimize the `depthwise` function by leveraging the power of TVM.
+# TVM community proposes an embedded Domain Specific Language on Python call TVMscript,
+# which serves for an abstraction of program on various hardware backends.


I think calling TVMScript as "an abstraction of program on various hardware backends" is a bit long shot. I think it is a much more high-level, concrete thing.

masahi · 2022-09-01T01:56:35Z

gallery/how_to/work_with_pytorch/using_as_torch.py

+# The computations and machine learning compilation analysis will be defined around them.
+# The last 3 lines are computation statements, including an initialization of `C[vj, vi]` and the summing up along the axis k.
+# Finally, we place 2 decorators `T.prim_func` and `as_torch` above the definition of function,
+# which converts the Python AST to TVMscript AST and then converts to PyTorch's `nn.Module`.


These sentences might be too detailed for a tutorial intended for PT users. I prefer a more succinct summary of what TVMScript is about, not necessarily explaining all the syntactic constructs used in the example.

masahi · 2022-09-01T01:59:13Z

gallery/how_to/work_with_pytorch/using_optimized_torch.py

+======================
+**Author**: `Yaoda Zhou <https://github.com/juda/>`_
+This article is an introductory tutorial to optimize PyTorch models by using `tvm.contrib.torch.optimize_torch`.
+For us to follow this tutorial, PyTorch, as well as TorchVision, should be installed.


I think you copied "For us to follow this tutorial" from other tutorials, but this is not a good English phrase. We can just say "To follow this tutorial".

masahi · 2022-09-01T02:02:25Z

gallery/how_to/work_with_pytorch/using_optimized_torch.py

+# Optimized SimpleModel by TVM MetaSchedule
+# ------------------------------
+# We provide a `optimize_torch` function, which has the similar usage as `torch.jit.trace`.
+# The optimized function/model and example input are required to provide by users.


The PyTorch model to optimize, along with its example input, are provided by users.

masahi · 2022-09-01T02:26:31Z

gallery/how_to/work_with_pytorch/using_optimized_torch.py

+# ------------------------------
+# Besides, let us define a resnet18 model in a standard way.
+# TorchScript also provides a built-in "optimize_for_inference" function to accelerate the inference,
+# we will compare the performance of those two optimizers later.


Same comment as "Define the resnet18 optimized by MetaSchedule" above.

masahi · 2022-09-01T02:26:55Z

gallery/how_to/work_with_pytorch/using_optimized_torch.py

+# we will compare the performance of those two optimizers later.
+
+
+class JitModule(torch.nn.Module):


Drop JitModule boilerplate.

Please address this comment. There is no need to have JitModule.

masahi · 2022-09-01T02:27:47Z

gallery/how_to/work_with_pytorch/using_optimized_torch.py

+jit_module_resnet18 = JitModule()
+
+######################################################################
+# Compare the performance between two scheduling approaches.


What are "two scheduling approaches"? torch.jit.optimize_for_inference is not a scheduling approach.

masahi · 2022-09-01T02:28:08Z

gallery/how_to/work_with_pytorch/using_optimized_torch.py

+        ).blocked_autorange()
+    )
+
+# We can print the results on screen.


Drop this sentence

masahi · 2022-09-01T02:29:04Z

gallery/how_to/work_with_pytorch/using_optimized_torch.py

+
+# In the working machine, the average inference time by `optimized_torch` is 860.5 us,
+# while the average inference time of `jit_optimized` is 1156.3 us,
+# showing the performance arises by around 1/4.


Apply my comment from as_torch tutorial here too. I won't repeat the same comment.

juda · 2022-09-08T03:14:18Z

@tvm-bot rerun

juda · 2022-09-09T07:05:33Z

@masahi I finish another round of polishing. Could you please have a look?

masahi · 2022-09-12T09:06:22Z

gallery/how_to/work_with_pytorch/using_as_torch.py

+======================
+**Author**: 
+`Yaoda Zhou <https://github.com/juda>`_,
+`Masahiro Masuda <https://github.com/masahi>`_


No need to add me as an author.

masahi · 2022-09-12T09:08:49Z

gallery/how_to/work_with_pytorch/using_as_torch.py

+compare = benchmark.Compare(results)
+compare.print()
+
+# In author's environment, the average inference time of `tvm_depthwise` is 120.0 us (TVM version is 0.9.0),


0.9.0 is the released version, I don't think this is the one you are using for development. There is no need to mention the TVM version.

masahi · 2022-09-13T07:13:37Z

gallery/how_to/work_with_pytorch/using_as_torch.py

+
+
+# Then, we plan to optimize the `depthwise` function by leveraging the power of TVM.
+# TVM community proposes an embedded Domain Specific Language on Python called TVMscript,


masahi · 2022-09-13T07:32:54Z

gallery/how_to/work_with_pytorch/using_as_torch.py

+
+# Then, we plan to optimize the `depthwise` function by leveraging the power of TVM.
+# TVM community proposes an embedded Domain Specific Language on Python called TVMscript,
+# serving for a high-level abstraction of TVM intermediate representative,


which serves as the high-level frontend for TVM's Tensor IR.

masahi · 2022-09-13T07:33:10Z

gallery/how_to/work_with_pytorch/using_as_torch.py

+# Then, we plan to optimize the `depthwise` function by leveraging the power of TVM.
+# TVM community proposes an embedded Domain Specific Language on Python called TVMscript,
+# serving for a high-level abstraction of TVM intermediate representative,
+# which is easy to impose transformations and optimizations and deploy on various hardware backends.


This sentence can be dropped

masahi · 2022-09-13T07:54:21Z

gallery/how_to/work_with_pytorch/using_optimized_torch.py

+# In such a way, we obtain a new resnet18 model optimized by MetaSchedule.
+
+
+class MyResNet18(torch.nn.Module):


Please address this comment. There is no need to have MyResNet18.

masahi · 2022-09-13T07:55:00Z

gallery/how_to/work_with_pytorch/using_optimized_torch.py

+# we will compare the performance of those two optimizers later.
+
+
+class JitModule(torch.nn.Module):


Please address this comment. There is no need to have JitModule.

masahi · 2022-09-13T07:55:25Z

gallery/how_to/work_with_pytorch/using_optimized_torch.py

+######################################################################
+# Compare the performance between two approaches.
+# ------------------------------
+# Using PyTorch's benchmark Compare class, we can have a direct comparison result between two inference models.


Drop this sentence.

masahi · 2022-09-13T07:57:27Z

gallery/how_to/work_with_pytorch/using_optimized_torch.py

+compare = benchmark.Compare(results)
+compare.print()
+
+# In author's environment, the average inference time of `tvm_module_resnet18` is 620.0 us (TVM version is 0.9.0),


Drop the reference to TVM version (see the same comment to using_as_torch.py)

masahi · 2022-09-13T07:58:06Z

gallery/how_to/work_with_pytorch/using_as_torch.py

+######################################################################
+# Benchmark
+# -------------------------------
+# We will compare two operators by using PyTorch's benchmark toolkit.


Drop this sentence. It is not useful.

masahi · 2022-09-13T08:02:42Z

gallery/how_to/work_with_pytorch/using_as_torch.py

+# specific language governing permissions and limitations
+# under the License.
+"""
+Wrap Your TVMscript with PyTorch Module


Is "Wrap ... with " the right phrase? I think "Wrap ... as PyTorch Module" is more correct.

masahi · 2022-09-13T08:03:28Z

gallery/how_to/work_with_pytorch/using_as_torch.py

+`Yaoda Zhou <https://github.com/juda>`_,
+`Masahiro Masuda <https://github.com/masahi>`_
+
+This article is an introductory tutorial on wrapping the TVMscript code with the PyTorch module.


This article is a tutorial on wrapping TVMScript code as a PyTorch module.

masahi · 2022-09-13T08:03:56Z

gallery/how_to/work_with_pytorch/using_as_torch.py

+`Masahiro Masuda <https://github.com/masahi>`_
+
+This article is an introductory tutorial on wrapping the TVMscript code with the PyTorch module.
+By the decorator `as_torch`, users can wrap a TVMscript code into a PyTorch nn.Module naturally.


"Using the decorator..."

Drop "a" before TVMScript

masahi · 2022-09-13T08:04:40Z

gallery/how_to/work_with_pytorch/using_optimized_torch.py

+`Yaoda Zhou <https://github.com/juda>`_,
+`Masahiro Masuda <https://github.com/masahi>`_
+
+This article is an introductory tutorial to optimize PyTorch models by using `tvm.contrib.torch.optimize_torch`.


an introductory tutorial -> a tutorial

juda · 2022-09-22T13:13:34Z

Hi @masahi , I polished the tutorial according to your feedback, could you please read it one more time?

masahi

This looks very good now.

masahi · 2022-09-27T06:22:14Z

@juda Thank your for your patience. I think the tutorial is now very clean and simple, without unnecessary things.

…ache#12318) * how-to use optmized_torch * as_torch * format * one more comment * improve doc * improve code * fix text * SSR * CPU model * whitespace * improve document * small edit * retrigger ci * using_as_torch polish * using_optimized_torch * fix errors * one more author * small edit * polish as_torch * save progress * more edit * small edit Co-authored-by: juda <yzhou@octoml.ai>

juda added 4 commits August 4, 2022 01:41

how-to use optmized_torch

d4ea333

as_torch

071ff9a

format

4b83b57

one more comment

b99264f

juda commented Aug 8, 2022

View reviewed changes

gallery/how_to/work_with_pytorch/using_as_torch.py Show resolved Hide resolved

juda added 7 commits August 8, 2022 19:42

improve doc

57bdf4b

improve code

bf98edc

fix text

e3289e4

SSR

7187943

Merge branch 'main' of github.com:juda/tvm into how-to

f4da268

CPU model

c1dc771

whitespace

e00fecc

masahi requested changes Aug 17, 2022

View reviewed changes

improve document

eaa2dfc

yelite reviewed Aug 29, 2022

View reviewed changes

juda added 3 commits August 30, 2022 20:41

small edit

388c7ba

retrigger ci

dca7dda

Merge branch 'main' of github.com:juda/tvm into how-to

9d13fc2

masahi requested changes Sep 1, 2022

View reviewed changes

juda added 4 commits September 4, 2022 20:03

using_as_torch polish

8d29add

Merge branch 'main' of github.com:juda/tvm into how-to

4ce8e3f

using_optimized_torch

3c2fd9e

fix errors

ca1331c

one more author

70b51ff

small edit

0430842

masahi requested changes Sep 13, 2022

View reviewed changes

masahi reviewed Sep 13, 2022

View reviewed changes

juda added 5 commits September 13, 2022 19:19

polish as_torch

de08328

save progress

88cc5cd

Merge branch 'main' of github.com:juda/tvm into how-to

325e119

more edit

88c3340

small edit

1ef893d

masahi approved these changes Sep 27, 2022

View reviewed changes

masahi merged commit b61f633 into apache:main Sep 27, 2022

leandron mentioned this pull request Feb 1, 2023

TVM v0.11.0 Release Candidate Notes #13899

Closed


		print(tvm_depthwise.script())

		# Hint: If user plan to deploy on GPU, the GPU target should be provided,

		# we will compare the performance of those two optimizers later.


		class JitModule(torch.nn.Module):



		# Then, we plan to optimize the `depthwise` function by leveraging the power of TVM.
		# TVM community proposes an embedded Domain Specific Language on Python called TVMscript,

		# In such a way, we obtain a new resnet18 model optimized by MetaSchedule.


		class MyResNet18(torch.nn.Module):

[TVM PyTorch Integration] optimized_torch & as_torch how-to guide #12318

[TVM PyTorch Integration] optimized_torch & as_torch how-to guide #12318

Conversation

juda commented Aug 5, 2022 • edited

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

juda commented Sep 1, 2022

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

juda commented Sep 8, 2022

juda commented Sep 9, 2022

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

juda commented Sep 22, 2022

masahi left a comment

Choose a reason for hiding this comment

masahi commented Sep 27, 2022

juda commented Aug 5, 2022 •

edited