@@ -21,6 +21,9 @@ For example, we can define a LeNet module like this:
2121.. code-block :: python
2222 :linenos:
2323
24+ import torch.nn as nn
25+ import torch.nn.functional as F
26+
2427 class LeNetFeatExtractor (nn .Module ):
2528 def __init__ (self
2629 super (LeNetFeatExtractor, self ).__init__ ()
@@ -56,6 +59,7 @@ For example, we can define a LeNet module like this:
5659 x = self .feat(x)
5760 x = self .classifer(x)
5861 return x
62+
5963
6064
6165 Obviously you may want to consolidate such a simple model into a single module but we can see the composability of PyTorch here
@@ -67,15 +71,20 @@ To trace an instance of our LeNet module, we can call ``torch.jit.trace`` with a
6771
6872.. code-block :: python
6973
74+ import torch.jit
75+
7076 model = LeNet()
71- traced_model = torch.jit.trace(model, torch.empty([1 ,1 ,32 ,32 ]))
77+ input_data = torch.empty([1 ,1 ,32 ,32 ])
78+ traced_model = torch.jit.trace(model, input_data)
7279
7380
7481is following the execution of your module, it cannot pick up control flow for instance. By working from the Python code, the compiler can
7582include these components. We can run the script compiler on our LeNet module by calling ``torch.jit.script ``
7683
7784.. code-block :: python
7885
86+ import torch.jit
87+
7988 model = LeNet()
8089 script_model = torch.jit.script(model)
8190
@@ -138,20 +147,23 @@ to load in a deployment application. In order to load a TensorRT/TorchScript mod
138147 import trtorch
139148
140149 ...
150+
151+ script_model.eval() # torch module needs to be in eval (not training) mode
152+
141153 compile_settings = {
142154 " input_shapes"
143155 {
144- " min" 1 , 3 , 224 , 224 ],
145- " opt" 1 , 3 , 512 , 512 ],
146- " max" 1 , 3 , 1024 , 1024 ]
147- }, # For static size [1, 3, 224, 224]
156+ " min" 1 , 1 , 32 , 32 ],
157+ " opt" 1 , 1 , 32 , 32 ],
158+ " max" 1 , 1 , 32 , 32 ]
159+ },
148160 ],
149- " op_precision" # Run with FP16
161+ " op_precision" # Run with fp16
150162 }
151163
152- trt_ts_module = trtorch.compile(torch_script_module , compile_settings)
164+ trt_ts_module = trtorch.compile(script_model , compile_settings)
153165
154- input_data = input_data.half()
166+ input_data = input_data.to( ' cuda ' ). half()
155167 result = trt_ts_module(input_data)
156168 torch.jit.save(trt_ts_module, " trt_ts_module.ts"
157169
@@ -162,7 +174,7 @@ to load in a deployment application. In order to load a TensorRT/TorchScript mod
162174 import trtorch
163175
164176 trt_ts_module = torch.jit.load(" trt_ts_module.ts"
165- input_data = input_data.half()
177+ input_data = input_data.to( ' cuda ' ). half()
166178 result = trt_ts_module(input_data)
167179
168180ts_in_cc :
0 commit comments