jit trace
torch.jit.trace使用eager model和一个dummy input作为输入,tracer会根据提供的model和input记录数据在模型中的流动过程,然后将整个模型转换为torchscript module。看一个具体的例子:
我们使用bert(bidirectional encoder representations from transformers)作为例子。
from transformers import berttokenizer, bertmodelimport numpy as npimport torchfrom time import perf_counterdef timer(f,*args): start = perf_counter() f(*args) return (1000 * (perf_counter() - start))# 加载bert modelnative_model = bertmodel.from_pretrained(bert-base-uncased)# huggingface的api中,使用torchscript=true参数可以直接加载torchscript modelscript_model = bertmodel.from_pretrained(bert-base-uncased, torchscript=true)script_tokenizer = berttokenizer.from_pretrained('bert-base-uncased', torchscript=true)# tokenizing input texttext = [cls] who was jim henson ? [sep] jim henson was a puppeteer [sep]tokenized_text = script_tokenizer.tokenize(text)# masking one of the input tokensmasked_index = 8tokenized_text[masked_index] = '[mask]'indexed_tokens = script_tokenizer.convert_tokens_to_ids(tokenized_text)segments_ids = [0, 0, 0, 0, 0, 0, 0, 1, 1, 1, 1, 1, 1, 1]# creating a dummy inputtokens_tensor = torch.tensor([indexed_tokens])segments_tensors = torch.tensor([segments_ids])然后分别在cpu和gpu上测试eager mode的pytorch推理速度。
# 在cpu上测试eager model推理性能native_model.eval()np.mean([timer(native_model,tokens_tensor,segments_tensors) for _ in range(100)])# 在gpu上测试eager model推理性能native_model = native_model.cuda()native_model.eval()tokens_tensor_gpu = tokens_tensor.cuda()segments_tensors_gpu = segments_tensors.cuda()np.mean([timer(native_model,tokens_tensor_gpu,segments_tensors_gpu) for _ in range(100)])再分别在cpu和gpu上测试script mode的torchscript模型的推理速度
# 在cpu上测试torchscript性能traced_model = torch.jit.trace(script_model, [tokens_tensor, segments_tensors])# 因模型的trace时,已经包含了.eval()的行为,因此不必再去显式调用model.eval()np.mean([timer(traced_model,tokens_tensor,segments_tensors) for _ in range(100)])# 在gpu上测试torchscript的性能最终运行结果如表
我使用的硬件规格是google colab,cpu是intel(r) xeon(r) cpu @ 2.00ghz,gpu是tesla t4。
从结果来看,在cpu上,torchscript比pytorch eager快了3.5%,在gpu上,torchscript比pytorch快了55.6%。
然后我们再用resnet做一个测试。
import torchvisionimport torchfrom time import perf_counterimport numpy as npdef timer(f,*args): start = perf_counter() f(*args) return (1000 * (perf_counter() - start))# pytorch cpu versionmodel_ft = torchvision.models.resnet18(pretrained=true)model_ft.eval()x_ft = torch.rand(1,3, 224,224)print(f'pytorch cpu: {np.mean([timer(model_ft,x_ft) for _ in range(10)])}')# pytorch gpu versionmodel_ft_gpu = torchvision.models.resnet18(pretrained=true).cuda()x_ft_gpu = x_ft.cuda()model_ft_gpu.eval()print(f'pytorch gpu: {np.mean([timer(model_ft_gpu,x_ft_gpu) for _ in range(10)])}')# torchscript cpu versionscript_cell = torch.jit.script(model_ft, (x_ft))print(f'torchscript cpu: {np.mean([timer(script_cell,x_ft) for _ in range(10)])}')# torchscript gpu versionscript_cell_gpu = torch.jit.script(model_ft_gpu, (x_ft_gpu))print(f'torchscript gpu: {np.mean([timer(script_cell_gpu,x_ft.cuda()) for _ in range(100)])}')
torchscript相比pytorch eager model,cpu性能提升4.2%,gpu性能提升45%。与bert的结论一致。
荣耀9什么时候上市?荣耀9发布在即:荣耀9配置,价格,外观消息汇总:麒麟960+3D曲面玻璃机身,下午2点半发布
宝马X5 xDrive30d上手体验评测
前景巨大的AoX蓝牙测向技术
微软否认推Linux版Office 称仅研究兼容性
谷歌首次出现集体请辞 锤子科技捐赠480万给OpenSSL与OpenBSD
TorchScript model与eager model的性能区别
浅析三星的晶圆代工梦
多层感知器的人工神经网络解析
工业以太网串口网关ES-301A/B的特点以及技术特性
数据科学家和数据工程师能合二为一吗?
小米不在MWC期间发布小米6 全为搭载自家松果处理器铺路!
协作机器人领域最大单笔融资花落节卡
手机探测电路图
关于高效检漏问题,他们有这些方法
势不可挡的医疗机器人,医生会被人工智能取代吗?
无功补偿常见问题有哪些
利用传感器技术和马达控制实现的工业自动化系统方案
未来传感器将向两大方面突破,五个方向发展
涡街流量计安装注意事项
vivo新机获得3C认证,支持120W快充,骁龙888旗舰