Package: towhee.compiler
Towhee CompilerTowhee compiler is a Python JIT compiler that speeds up AI-related codes by native code generation. The project is inspired byNumba,PyjionandTorchDynamo. Towhee compiler uses a frame evaluation hook (see [PEP 523]:https://www.python.org/dev/peps/pep-0523/) to get the chance of compiling python bytecodes into native code.The code is based on a forked version of torchdynamo, which extractfx.Graphby trace the execution of python code. But the goal of towhee compiler iswhole program code generation, which also includes program that can not be represented byfx.Graph.InstallInstall with pipIf install towhee.compiler failed with pip, please install it from source.$pipinstalltowhee.compilerInstall from source code$gitclone-b0.1.0https://github.com/towhee-io/towhee-compiler.git
$cdtowhee-compiler&&pipinstall-rrequirements
$python3setup.pydevelopExamplesRun with Torch ModelCompileTowhee compiler can speedup any models, for example, we just need to addjit_compilecontext to theimage_embeddingfunction.importtorchimporttorchvision.modelsasmodelsimportnumpyasnpimporttowhee.compilerfromtowhee.compilerimportjit_compile# towhee.compiler.config.debug = Truetorch_model=models.resnet50()torch_model=torch.nn.Sequential(*(list(torch_model.children())[:-1]))torch_model=torch_model.eval()defimage_embedding(inputs):imgs=torch.tensor(inputs)embedding=torch_model(imgs).detach().numpy()returnembedding.reshape([2048])inputs=np.random.randn(1,3,244,244).astype(np.float32)withjit_compile():embeddings=image_embedding(inputs)TimerWe have compiled the model with the nebullvm backend (the default backend in towhee.compiler ), and we can define a Timer class to record the time spent.importtimeclassTimer:def__init__(self,name):self._name=namedef__enter__(self):self._start=time.time()returnselfdef__exit__(self,*args):self._interval=time.time()-self._startprint('%s:%.2fs'%(self._name,self._interval))And we can see that the compiled function is more than 3 times faster.withTimer('Image Embedding'):embeddings=image_embedding(inputs)withTimer('Image Embedding with towhee compiler'),jit_compile():embeddings_jit=image_embedding(inputs)Image Embedding: 0.14sImage Embedding with towhee compiler: 0.04sRun with TowheeTowhee supports setting JIT to usetowhee.compilerto compile.Set JITFor example, we can addset_jit('towhee')in image embedding pipeline, then the following operator will be automatically compiledimporttowheeembeddings_towhee=(towhee.dc(['https://raw.githubusercontent.com/towhee-io/towhee/main/towhee_logo.png']).image_decode().set_jit('towhee').image_embedding.timm(model_name='resnet50'))TimerAnd we can make two towhee pipeline function to record the time cost.towhee_func=(towhee.dummy_input().image_embedding.timm(model_name='resnet50').as_function())towhee_func_jit=(towhee.dummy_input().set_jit('towhee').image_embedding.timm(model_name='resnet50').as_function())data=towhee.ops.image_decode()('https://raw.githubusercontent.com/towhee-io/towhee/main/towhee_logo.png')withTimer('Towhee function'):emb=towhee_func(data)withTimer('Towhee function with Compiler'):emb_jit=towhee_func_jit(data)Towhee function: 0.14sTowhee function with Compiler: 0.08sTests in Towhee HubAccording to the README of Operator onTowhee Hub, we set jit to compile and speedup model , theresults are as follows:5.5 means that the performance after jit is 5.5 times, and N means no speedup or compilation failure. And more test results will be updated continuously.FieldTaskOperatorSpeedup(CPU/GPU)ImageImage Embeddingimage_embedding.timm1.3/1.3image_embedding.data2vec1.2/1.7image_embedding.swag1.4/NFace Embeddingface_embedding.inceptionresnetv13.2/NFace Landmarkface_landmark_detection.mobilefacenet2.1/2.1NLPText Embeddingtext_embedding.transformers2.6/Ntext_embedding.data2vec1.8/Ntext_embedding.realm5.5/1.9text_embedding.xlm_prophetnet2.1/2.8AudioAudio Classificationaudio_classification.panns1.6/NAudio Embeddingaudio_embedding.vggish1.5/Naudio_embedding.data2vec1.5/NMultimodalImage Textimage_text_embedding.blip2.3/NVideo Textvideo_text_embedding.bridge_former(modality='text')2.1/Nvideo_text_embedding.frozen_in_time(modality='text')2.2/N