Package: rtx-deep
rtx_deep: Deep AI modules developed by MOGO RTX team, aims to accelerate the distributed training, int8-aware distributed training, distributed evaluation and inference, model tracing and optimization, and TensorRT deployment.1 Dependencytorch>=1.10.0
tensorrt>=7.0
graphviz2 Installationpip3installgraphviz
apt-getinstallgraphviz
python3setup.pyinstall3 Examples3.1 Graph Tracing and Model Optimizationimporttorch
importtorch.nnasnn
importtorch.nn.functionalasF

importrtx_deep


classconv3x3_bn_relu(nn.Module):def__init__(self,in_planes,out_planes,stride=1,dilation=1,groups=1):super(conv3x3_bn_relu,self).__init__()self.net=nn.Sequential(nn.Conv2d(in_planes,out_planes,kernel_size=3,stride=stride,padding=dilation,dilation=dilation,groups=groups,bias=False),nn.BatchNorm2d(out_planes),nn.ReLU(inplace=True))defforward(self,x):x1=self.net(x)returnx1


classModel(nn.Module):def__init__(self):super(Model,self).__init__()self.net=nn.Sequential(conv3x3_bn_relu(64,64),conv3x3_bn_relu(64,64))defforward(self,x):x1=self.net(x)returnx1model=Model()model.eval()model.cuda()input_data=torch.randn(1,64,1024,1024).cuda()# graph tracingmodel_fx=rtx_deep.graph_tracer.ad_trace.graph_trace(model,function_name=None)# Model Optimization# conduct graph tracing in graph_optim_from_module automaticallymodel_fx_optim=rtx_deep.graph_tracer.graph_utils.graph_optim_from_module(model,function_name=None,sample_inputs=(input_data,))3.2 Quantization-Aware Trainingimporttorch
importtorch.nnasnn
importtorch.nn.functionalasF

importrtx_deep


classconv3x3_bn_relu(nn.Module):def__init__(self,in_planes,out_planes,stride=1,dilation=1,groups=1):super(conv3x3_bn_relu,self).__init__()self.net=nn.Sequential(nn.Conv2d(in_planes,out_planes,kernel_size=3,stride=stride,padding=dilation,dilation=dilation,groups=groups,bias=False),nn.BatchNorm2d(out_planes),nn.ReLU(inplace=True))defforward(self,x):x1=self.net(x)returnx1


classModel(nn.Module):def__init__(self):super(Model,self).__init__()self.net=nn.Sequential(conv3x3_bn_relu(64,64),conv3x3_bn_relu(64,64))defforward(self,x):x1=self.net(x)returnx1model=Model()model.eval()model.cuda()input_data=torch.randn(1,64,1024,1024).cuda()# Model Optimization# conduct graph tracing in graph_optim_from_module automaticallymodel_fx_optim=rtx_deep.graph_tracer.graph_utils.graph_optim_from_module(model,function_name=None)# qatmodel_qat=rtx_deep.quant_lib.quant_utils.prepare_qat(model_fx_optim,sample_inputs=[input_data],observe_config_dic=dict(averaging_constant=0.05),quant_config_dic=dict(quant_min=-127,quant_max=127,is_symmetric=True,is_quant=True),disable_prefix=[])# vis model networkrtx_deep.graph_tracer.vis_model.vis(model_fx_optim,'./model_fx_optim.png')rtx_deep.graph_tracer.vis_model.vis(model_qat,'./model_qat.png')# qat training...3.3 TensorRT Deploymentimporttorch
importtorch.nnasnn
importtorch.nn.functionalasF

importrtx_deep
importrtx_deep_plugin
fromrtx_deep.deploy_lib.convert_trtimportInputTensor,torch2trt

classconv3x3_bn_relu(nn.Module):def__init__(self,in_planes,out_planes,stride=1,dilation=1,groups=1):super(conv3x3_bn_relu,self).__init__()self.net=nn.Sequential(nn.Conv2d(in_planes,out_planes,kernel_size=3,stride=stride,padding=dilation,dilation=dilation,groups=groups,bias=False),nn.BatchNorm2d(out_planes),nn.ReLU(inplace=True))defforward(self,x):x1=self.net(x)returnx1


classModel(nn.Module):def__init__(self):super(Model,self).__init__()self.net=nn.Sequential(conv3x3_bn_relu(64,64),conv3x3_bn_relu(64,64))defforward(self,x):x1=self.net(x)x2=rtx_deep_plugin.max_op(x1,dim=1)returnx2model=Model()model.eval()model.cuda()input_data=torch.randn(1,64,1024,1024).cuda()# Model Optimization# conduct graph tracing in graph_optim_from_module automaticallymodel_fx_optim=rtx_deep.graph_tracer.graph_utils.graph_optim_from_module(model,function_name=None,sample_inputs=(input_data,))# TensorRT Deploymentmodel_trt=torch2trt(model=model_fx_optim,input_specs=[InputTensor(input_data,'input_data')],output_names=['max_value','max_index'],fp16_mode=True,#dla_core=0,strict_type_constraints=True,explicit_precision=True)# vis tensorrt networkrtx_deep.deploy_lib.tools.vis_trt.vis(model_trt.network,'test.png')error=model(input_data)[0]-model_trt(input_data)[0]print(error.abs().max())