EDUCBA Logo

EDUCBA

MENUMENU
  • Explore
    • EDUCBA Pro
    • PRO Bundles
    • Featured Skills
    • New & Trending
    • Fresh Entries
    • Finance
    • Data Science
    • Programming and Dev
    • Excel
    • Marketing
    • HR
    • PDP
    • VFX and Design
    • Project Management
    • Exam Prep
    • All Courses
  • Blog
  • Enterprise
  • Free Courses
  • Log in
  • Sign Up
Home Software Development Software Development Tutorials PyTorch Tutorial PyTorch JIT
 

PyTorch JIT

Updated April 5, 2023

PyTorch JIT

 

 

Introduction to PyTorch JIT

PyTorch JIT is used for compiling the programs written in PyTorch. It is a completely optimized compiler solution for PyTorch. PyTorch JIT comes with several features and characteristics that make it perfect. Some of them include having the additional feature of interpreting in a lightweight and safe environment, custom transformations can be written easily and supportive, and it comes along with automated diff support along with the inferencing.

Watch our Demo Courses and Videos

Valuation, Hadoop, Excel, Mobile Apps, Web Development & many more.

In this article, we will learn about what PyTorch JIT is, using the PyTorch JIT, JIT Trace, and JIT Script, PyTorch JIT Modules, some of the examples associated with it, and finally, conclude our statement.

What is PyTorch JIT?

JIT compiler makes use of runtime information for the module optimization of torchscript. Using JIT that is Just In Time compiler automation in optimizing sparsification, quantization, and layer fusion are possible.

Script mode makes the use of eager module and torch.JIT.script and torch.JIT.trace modules for creating an IR that is an Intermediate Representation of PyTorch.

Using the PyTorch JIT

We need to create models, transfer them to the production environment, and then compile and optimize it using the PyTorch JIT compiler, making inferencing faster without much effort.

PyTorch handles production and research using two modes: eager and script mode. Eager is used for training, prototyping, and experimentation, while script mode is for production. Script mode contains two separate components, namely TorchScript and PyTorch JIT. Script mode proves to be of great help to avoid dependence on python and python’s GIL. Moreover, it brings along features like portability and performance due to the use of PyTorch JIT, which is an optimized compiler.

JIT Trace and JIT Script

The invocation of script mode is carried out either by using a torch.JIT.script or by using a torch.JIT.trace. JIT trace that is a torch.JIT.trace accepts the input containing our trained, eager model and the instance of data. The JIT tracer is responsible for keeping the tensor’s operations records and helps run the supplied module. Further, we convert the recordings made by the JIT tracer to the module of the torch script.

The JIT script is also known as a torch.JIT.script makes us write the code straight inside the torch script. It is versatile and, at the same time, a little verbose. After certain adjustments, it helps us support most of the PyTorch models. We just need to provide the model instance to the JIT script, which is completely different than that of the JIT trace mode. It is not necessary to specify the sample of data.

PyTorch JIT Modules

Script mode makes the use of eager module and torch.JIT.script and torch.JIT.trace modules for creating an IR that is an Intermediate Representation of PyTorch. PyTorch JIT modules are the JIT trace module and JIT script module.

JIT Trace Module

Tracer can reuse the coder if the eager model exists and handle all the programs using exclusive torch operations and tensors. The disadvantages of tracer include omitting the python constructs, data structures, and flow of control. Even without specification of any warnings, it creates representations they are not faithful. Just to verify that the structures of the model of PyTorch are being parsed correctly, it always checks the IR.

JIT Script Module

It looks similar to python and supports the python constructs, and the flow of control is preserved. In addition, it provides complete support for dictionaries and lists. However, it does not support the constant values, and typecasting is necessary. The default type, when not specified anything, is tensor.

Example of PyTorch JIT

Let us consider one example that was developed by google artificial intelligence named BERT that stands for Bidirectional Encoder Representations from Transformers. The library that we will use is provided by hugging face.

We will follow the steps given below to implement the example –

  • Tokenizers and the BERT model will be initialized, and sample data will be created for inference.
  • PyTorch models will be prepared for inferencing on both GPU and CPU.
  • Torchscript models should be prepared for inferencing on both GPU and CPU.
from transformers
import numpy as numpyObj
import torch
import BertTokenizer, BertModel
from time import perf_counter
def sampleTimerFun(f,*args):
initiate = perf_counter()
f(*args)
return (1000 * (perf_counter() - initiate))
sampleTokenizerForScript = BertTokenizer.from_pretrained('bert-base-uncased', torchscript=True)
modelOfScript = BertModel.from_pretrained("bert-base-uncased", torchscript=True)
# Text in input is being tokenized
inputString = "[CLS] Who was Jim Henson ? [SEP] Jim Henson was a puppeteer [SEP]"
tokenized_inputString = sampleTokenizerForScript.tokenize(inputString)
# One of the input token is being masked
indexOfMaskedToken = 8
tokenized_inputString[indexOfMaskedToken] = '[MASK]'
ListOfIndexedTokens = sampleTokenizerForScript.convert_tokens_to_ids(tokenized_inputString)
IdOfSegments = [0, 0, 0, 0, 0, 0, 0, 1, 1, 1, 1, 1, 1, 1]
# dummy object creation for input value
tensorOfTokens = torch.tensor([ListOfIndexedTokens])
tensorOfSegments = torch.tensor([IdOfSegments])

# Example 1.1 – Sample Example demonstrating BERT model on CPU

sampleNativeModel = BertModel.from_pretrained("bert-base-uncased")
numpyObj.mean([sampleTimerFun(sampleNativeModel,tensorOfTokens,tensorOfSegments) for _ in range(100)])

# Example 1.2 – Sample Example demonstrating BERT model on GPU

# In order for inference to be applied, both of the models used for sampling should be present on the sample GPU device.

sampleGPUOfNative = sampleNativeModel.cuda()
tensorOfTokens_gpu = tensorOfTokens.cuda()
tensorOfSegments_gpu = tensorOfSegments.cuda()
numpyObj.mean([sampleTimerFun(sampleGPUOfNative,tensorOfTokens_gpu,tensorOfSegments_gpu) for _ in range(100)])

#  Example 1.3 –  CPU implementation of the torch.jit.trace

JITTraceModel = torch.jit.trace(modelOfScript, [tensorOfTokens, tensorOfSegments])
numpyObj.mean([sampleTimerFun(JITTraceModel,tensorOfTokens,tensorOfSegments) for _ in range(100)])

# Example 1.4 – GPU implementation of the torch.jit.trace

JITTraceModel_gpu = torch.jit.trace(modelOfScript.cuda(), [tensorOfTokens.cuda(), tensorOfSegments.cuda()])
numpyObj.mean([sampleTimerFun(JITTraceModel_gpu,tensorOfTokens.cuda(),tensorOfSegments.cuda()) for _ in range(100)])

The execution of various examples above on CPU, GPU for trace mode, and BERT mode in PyTorch and torchscript models give the following inference output –

PyTorch JIT output 1

We can observe that runtimes are almost the same on CPU, but on GPU, the runtime of torchscript proves to be better than PyTorch.

Internally, torchcript does the creation of an Intermediate representation of the PyTorch models that can be optimistically compiled at runtime by using the Just In Time (JIT) compilers. We can have a look at the IR that is an intermediate representation by using the traced_model.code.

Conclusion

We learned about the script mode and PyTorch JIT. Order to create the models for research in a production environment can be created by using script mode, and at the same time, it remains in the ecosystem of PyTorch.

Recommended Articles

We hope that this EDUCBA information on “PyTorch JIT” was beneficial to you. You can view EDUCBA’s recommended articles for more information.

  1. Dataset Pytorch
  2. PyTorch Versions
  3. What is PyTorch?
  4. PyTorch TensorBoard

Primary Sidebar

Footer

Follow us!
  • EDUCBA FacebookEDUCBA TwitterEDUCBA LinkedINEDUCBA Instagram
  • EDUCBA YoutubeEDUCBA CourseraEDUCBA Udemy
APPS
EDUCBA Android AppEDUCBA iOS App
Blog
  • Blog
  • Free Tutorials
  • About us
  • Contact us
  • Log in
Courses
  • Enterprise Solutions
  • Free Courses
  • Explore Programs
  • All Courses
  • All in One Bundles
  • Sign up
Email
  • [email protected]

ISO 10004:2018 & ISO 9001:2015 Certified

© 2025 - EDUCBA. ALL RIGHTS RESERVED. THE CERTIFICATION NAMES ARE THE TRADEMARKS OF THEIR RESPECTIVE OWNERS.

EDUCBA

*Please provide your correct email id. Login details for this Free course will be emailed to you
Loading . . .
Quiz
Question:

Answer:

Quiz Result
Total QuestionsCorrect AnswersWrong AnswersPercentage

Explore 1000+ varieties of Mock tests View more

EDUCBA

*Please provide your correct email id. Login details for this Free course will be emailed to you
EDUCBA
Free Software Development Course

Web development, programming languages, Software testing & others

By continuing above step, you agree to our Terms of Use and Privacy Policy.
*Please provide your correct email id. Login details for this Free course will be emailed to you
EDUCBA

*Please provide your correct email id. Login details for this Free course will be emailed to you

EDUCBA Login

Forgot Password?

🚀 Limited Time Offer! - 🎁 ENROLL NOW