@@ -138,21 +138,18 @@ def init_model():
138
138
# Note that in the call to ``torch.compile``, we have have the additional
139
139
# ``mode`` argument, which we will discuss below.
140
140
141
- def evaluate (mod , inp ):
142
- return mod (inp )
143
-
144
141
model = init_model ()
145
142
146
143
# Reset since we are using a different mode.
147
144
import torch ._dynamo
148
145
torch ._dynamo .reset ()
149
146
150
- evaluate_opt = torch .compile (evaluate , mode = "reduce-overhead" )
147
+ model_opt = torch .compile (model , mode = "reduce-overhead" )
151
148
152
149
inp = generate_data (16 )[0 ]
153
150
with torch .no_grad ():
154
- print ("eager:" , timed (lambda : evaluate ( model , inp ))[1 ])
155
- print ("compile:" , timed (lambda : evaluate_opt ( model , inp ))[1 ])
151
+ print ("eager:" , timed (lambda : model ( inp ))[1 ])
152
+ print ("compile:" , timed (lambda : model_opt ( inp ))[1 ])
156
153
157
154
######################################################################
158
155
# Notice that ``torch.compile`` takes a lot longer to complete
@@ -166,7 +163,7 @@ def evaluate(mod, inp):
166
163
for i in range (N_ITERS ):
167
164
inp = generate_data (16 )[0 ]
168
165
with torch .no_grad ():
169
- _ , eager_time = timed (lambda : evaluate ( model , inp ))
166
+ _ , eager_time = timed (lambda : model ( inp ))
170
167
eager_times .append (eager_time )
171
168
print (f"eager eval time { i } : { eager_time } " )
172
169
@@ -176,7 +173,7 @@ def evaluate(mod, inp):
176
173
for i in range (N_ITERS ):
177
174
inp = generate_data (16 )[0 ]
178
175
with torch .no_grad ():
179
- _ , compile_time = timed (lambda : evaluate_opt ( model , inp ))
176
+ _ , compile_time = timed (lambda : model_opt ( inp ))
180
177
compile_times .append (compile_time )
181
178
print (f"compile eval time { i } : { compile_time } " )
182
179
print ("~" * 10 )
@@ -250,6 +247,10 @@ def train(mod, data):
250
247
# Again, we can see that ``torch.compile`` takes longer in the first
251
248
# iteration, as it must compile the model, but in subsequent iterations, we see
252
249
# significant speedups compared to eager.
250
+ #
251
+ # We remark that the speedup numbers presented in this tutorial are for
252
+ # demonstration purposes only. Official speedup values can be seen at the
253
+ # `TorchInductor performance dashboard <https://hud.pytorch.org/benchmark/compilers>`__.
253
254
254
255
######################################################################
255
256
# Comparison to TorchScript and FX Tracing
0 commit comments