Package: nobuco
diliajlNoBullshitConverter is a tool that helps you translate Pytorch models into Keras/Tensorflow graphs without losing your mind.Supports a wide range of architecturesControl flow ops (If, For, While)Recurrent layers (LSTM, GRU)Arbitrary torch functionsSimpleFlexibleEfficientSanity-preserving, with clear mistake messagingInstallationpipinstall-UnobucoTable of ContentsEssentialsChannel order wizardryIn-place operationsImplementation mismatch: pick your poisonGoing dynamicControl flowsDynamic shapesAd hoc modificationsSo we put a converter inside your converterBut was it worth it?Nobuco knowledge baseEssentialsSuppose we want to convert a Pytorch module similar to this one:classMyModule(nn.Module):def__init__(self):super().__init__()self.conv=nn.Conv2d(3,16,kernel_size=(3,3),padding=(1,1),stride=(2,2))defforward(self,x):x=self.conv(x)x=nn.Hardsigmoid()(x)x=1-x[:,::2]*x[:,1::2]returnxThe process is exactly what you would expect. Instantiate the module, create dummy inputs and call the magic function:importnobucofromnobucoimportChannelOrder,ChannelOrderingStrategyfromnobuco.layers.weightimportWeightLayerdummy_image=torch.rand(size=(1,3,256,256))pytorch_module=MyModule().eval()keras_model=nobuco.pytorch_to_keras(pytorch_module,args=[dummy_image],kwargs=None,inputs_channel_order=ChannelOrder.TENSORFLOW,outputs_channel_order=ChannelOrder.TENSORFLOW)Aaaand done! That's all it takes to... hold on, what's that?Nobuco says it doesn't know how to handle hard sigmoid.
Apparently, it's our job to provide a node converter for eitherF.hardsigmoidor the enclosingHardsigmoidmodule (or the entireMyModule, but that makes little sense). Here, we'll go for the former.Conversion is done directly. No layers upon layers of abstraction, no obscure intermediate representation.
A node converter is just aCallablethat takes in the same arguments as the corresponding node in Pytorch and outputs an equivalent node in Tensorflow.
The converted node preserves the original node's signature, but Pytorch tensors replaced with Tensorflow counterparts (be thattf.Tensor,KerasTensor,tf.Variable, orResourceVariable).This should do the trick:@nobuco.converter(F.hardsigmoid,channel_ordering_strategy=ChannelOrderingStrategy.MINIMUM_TRANSPOSITIONS)defhardsigmoid(input:torch.Tensor,inplace:bool=False):returnlambdainput,inplace=False:tf.keras.activations.hard_sigmoid(input)It works, but the outputs don't quite match. Perhaps we should check howPytorchandTensorflowdefine hard sigmoid.
And sure enough, their implementations differ. Have to type in the formula manually, I guess...@nobuco.converter(F.hardsigmoid,channel_ordering_strategy=ChannelOrderingStrategy.MINIMUM_TRANSPOSITIONS)defhardsigmoid(input:torch.Tensor,inplace:bool=False):returnlambdainput,inplace=False:tf.clip_by_value(input/6+1/2,clip_value_min=0,clip_value_max=1)And the happy result:The example above is artificial, but it illustrates the point.
It's not feasible to provide a node converter for every existing Pytorch op. There's literallythousandsof them!
Best we can do without the converter constantly lacking essential functionality, being riddled with bugs, doing weird stuff and breaking apart with every other PT/TF release
is to keep the tool simple and customizable, make it clear where a problem comes from and let theusersort things out.
Usually it's easy for a human to translate an isolated operation from one framework to another.
Reproducing the graph structure is a different matter entirely. For that, Nobuco has you covered!Nobuco lets you intervene in conversion at each step, asks for help where needed and doesn't bother you with routine stuff.https://user-images.githubusercontent.com/2457934/233740603-cc11acc5-cd6b-48c8-b089-ff3ead772dd0.mp4With an IDE, you can jump right where the node was [I]nvoked, [D]efined and [C]onvertedChannel order wizardrySome operations assume its input tensors have a channel dimension.
And as you probably know, Pytorch and Tensorflow do not agree on the layout of such tensors.
Pytorch adopts channel-first layout (BCH,BCHW, etc.)
while Tensorflow works efficiently with channel-last tensors (BHC,BHWC, ...).
Transposing tensors between the two layouts incurs non-trivial overhead as generally, tensor data must be physically rearranged.
In an effort to keep that overhead to the minimum, Nobuco does layout coercionslazily.
A couple of things are needed to make it possible:Tensorflow tensors are augmented with an additional property which stores their channel order, either Pytorch (channel first) or Tensorflow (channel last) style.Node converters have requirements on what channel order their inputs must have. Said requirements are expressed withchannel_ordering_strategyargument.Channel ordering strategies areFORCE_TENSORFLOW_ORDERInput tensors will be coerced to Tensorflow channel order.Convenient for converting channel-aware operations (convolution, batchnorm).FORCE_PYTORCH_ORDERInput tensors entering the node will look exactly as they do in the original Pytorch graph.Use it when the node does not interpret its input tensors as having a channel dimension (linear, matmul).MINIMUM_TRANSPOSITIONSThe channel order is decided by a majority vote (whichever prevails among the inputs). This way the number of coercions (i.e. tensor transpositions) is kept to the minimum.
It also means whenever there's only one input, it will be left untouched.Best choice for element-wise ops (most activations).MANUALYou are on your own. In exchange for unrestricted freedom, you take responsibility to coerce input tensors to suitable channel order and to also annotate output tensors with their order.The simple lazy approach makes wonders in most situations, but sometimes it produces suboptimal graphs.
Consider the code below. Imagine this is some sort of text processing network.
It first applies a GRU layer which assumes the inputs do not have a channel dimension, so its input/output layouts are the same in both Pytorch and Tensorflow.
But then, the outputs are passed to a couple of 1D convolutions which are channel-aware.
Because of that, a transpose op must be put in the converted graph.classMyModule(nn.Module):def__init__(self):super().__init__()self.gru=nn.GRU(32,128,num_layers=1,batch_first=True,bidirectional=False)self.conv1=nn.Conv1d(12,40,kernel_size=3,padding=1)self.conv2=nn.Conv1d(12,60,kernel_size=1,padding=0)defforward(self,x):x,hx=self.gru(x)x1=self.conv1(x)x2=self.conv2(x)returnx1,x2pytorch_module=MyModule().eval()inputs=[torch.normal(0,1,size=(1,12,32)),]keras_model=nobuco.pytorch_to_keras(pytorch_module,inputs,inputs_channel_order=ChannelOrder.PYTORCH,)The laziness shoots us in the foot here, and we get not one transpose but two:For such occasions, there's two brethren functions:force_tensorflow_orderandforce_pytorch_order.x,hx=self.gru(x)x=nobuco.force_tensorflow_order(x)x1=self.conv1(x)x2=self.conv2(x)In case you are curious, the implementation is trivial:@nobuco.traceabledefforce_tensorflow_order(inputs):returninputs@nobuco.converter(force_tensorflow_order,channel_ordering_strategy=ChannelOrderingStrategy.FORCE_TENSORFLOW_ORDER)defconverter_force_tensorflow_order(inputs):returnlambdainputs:inputsforce_pytorch_orderis defined analogously.In-place operationsNobuco can handle most situations where tensors are modified in-place. For instance, these will work just fine:classMyModule(nn.Module):defforward(self,x):x[:,1:2,16:25,8::2]*=2torch.relu_(x)returnxHowever, applying in-place operation to a slice yields incorrect result. What gives?classMyModule(nn.Module):defforward(self,x):torch.relu_(x[:,1:2,16:25,8::2])returnxYou see, Tensorflow graphs (and many other formats like ONNX) do not support in-place ops.
So when we take slice (x[:, 1:2, 16:25, 8::2]) in TF/ONNX, the result is not a view of the original tensor but a copy.
This copy is then passed torelu(which is not in-place either), and its result is not used anywhere.
As you can see above, the output tensors of__getitem__andrelu_aregrayed out, and these operations are excluded from the graph.
In fact, it's empty:The easiest way of fixing this is to explicitly assign the result to the slice.
Conveniently enough, most standard in-place operations in Pytorch do return their modified arguments as outputs.classMyModule(nn.Module):defforward(self,x):x[:,1:2,16:25,8::2]=torch.relu_(x[:,1:2,16:25,8::2])returnxImplementation mismatch: pick your poisonSometimes, Pytorch and Tensorflow just don't go along.
In case ofhardsigmoid, it's a mere inconvenience, but it can be much more sinister.Take the model below, for example.classMyModule(nn.Module):def__init__(self):super().__init__()self.factor=4self.conv=nn.Conv2d(3*self.factor**2,3*self.factor**2,kernel_size=1)defforward(self,x):x=nn.PixelUnshuffle(self.factor)(x)x=self.conv(x)x=nn.PixelShuffle(self.factor)(x)returnxIdeally, there would only be three nodes in the converted graph. That's not what we get, though.Tensorflow does not havepixel_unshuffle/pixel_shuffle.
Their closest counterparts,tf.nn.space_to_depth/tf.nn.depth_to_space,
do almost the same thing but not quite: output channels are in a different order.
The order must be fixed with a priceytranspose, no way around that. Or is there?Instead of emulating an absent Pytorch op in Tensorflow,
we might do the procedure in reverse: provide a Pytorch implementation for the Tensorflow node we want to convert to.
The overhead would be carried by the original Pytorch model leaving the converted graph nice and clean.fromnobuco.addons.torch.space_to_depthimportSpaceToDepthfromnobuco.addons.torch.depth_to_spaceimportDepthToSpaceclassMyModuleTFOptimized(nn.Module):def__init__(self):super().__init__()self.factor=4self.conv=nn.Conv2d(3*self.factor**2,3*self.factor**2,kernel_size=1)defforward(self,x):x=SpaceToDepth(self.factor)(x)x=self.conv(x)x=DepthToSpace(self.factor)(x)returnxTorch-optimizedTensorflow-optimizedTorch implementationF.pixel_unshuffleTensorflow converter@nobuco.converter(F.pixel_unshuffle,channel_ordering_strategy=ChannelOrderingStrategy.FORCE_TENSORFLOW_ORDER)defconverter_pixel_unshuffle(input:Tensor,downscale_factor:_int):deffunc(input,downscale_factor):x=tf.nn.space_to_depth(input,downscale_factor)x=channel_interleave2d(x,downscale_factor,reverse=True)returnxreturnfuncdefchannel_interleave2d(x,block_size:int,reverse:bool):b,h,w,c=x.shapen_blocks=block_size**2ifreverse:x=tf.reshape(x,(b,h,w,n_blocks,c//n_blocks))else:x=tf.reshape(x,(b,h,w,c//n_blocks,n_blocks))x=tf.transpose(x,(0,1,2,4,3))x=tf.reshape(x,(b,h,w,c))returnxTorch implementationclassSpaceToDepth(nn.Module):def__init__(self,block_size):super().__init__()self.block_size=block_sizedefforward(self,input):x=F.pixel_unshuffle(input,self.block_size)x=channel_interleave2d(x,self.block_size,reverse=False)returnxdefchannel_interleave2d(x:torch.Tensor,block_size:int,reverse:bool)->torch.Tensor:b,c,h,w=x.shapen_blocks=block_size**2ifreverse:x=x.view(b,n_blocks,c//n_blocks,h,w)else:x=x.view(b,c//n_blocks,n_blocks,h,w)x=x.transpose(1,2).reshape(b,c,h,w)returnxTensorflow converter@nobuco.converter(SpaceToDepth,channel_ordering_strategy=ChannelOrderingStrategy.FORCE_TENSORFLOW_ORDER)defconverter_space_to_depth(self,input:torch.Tensor):returnlambdainput:tf.nn.space_to_depth(input,self.block_size)Going dynamicControl flowsIntroducing python control flow statements into the compute graph is no easy feat.
Tensorflow can do so viatf.autograph, but at a cost ofsystem's complexityand with some notablelimitations.
Stuff like that is way above Nobuco's paygrade, so the following module cannot be properly handled without human intervention.classControlIf(nn.Module):def__init__(self):super().__init__()self.conv_pre=nn.Conv2d(3,16,kernel_size=(1,1))self.conv_true=nn.Conv2d(16,32,kernel_size=(1,1))self.conv_false=nn.Conv2d(16,32,kernel_size=(1,1))self.conv_shared=nn.Conv2d(32,32,kernel_size=(1,1))defforward(self,x):x=self.conv_pre(x)ifx.mean()>0:x=self.conv_true(x)x=torch.tanh(x)x=self.conv_shared(x)x=x+1else:x=self.conv_false(x)x=torch.sigmoid(x)x=self.conv_shared(x)x=x-1x=self.conv_shared(x)returnxOf course, it's possible to translate the dynamic module into a Tensorflow layer
(don't forget to decorate it with@tf.functionfor autograph to kick in).
But what if it contains inner modules, do you replicate them in Tensorflow all by hand?
Not unless you want to!
Just convert them separately and use the resulting graphs inside the parent layer.classControlIfKeras(tf.keras.layers.Layer):def__init__(self,conv_pre,conv_true,conv_false,conv_shared,*args,**kwargs):super().__init__(*args,**kwargs)self.conv_pre=conv_preself.conv_true=conv_trueself.conv_false=conv_falseself.conv_shared=conv_shareddefget_config(self):config=super().get_config()config.update({"conv_pre":self.conv_pre,"conv_true":self.conv_true,"conv_false":self.conv_false,"conv_shared":self.conv_shared,})returnconfig@tf.functiondefcall(self,x):x=self.conv_pre(x)iftf.reduce_mean(x)>0:x=self.conv_true(x)x=tf.tanh(x)x=self.conv_shared(x)x=x+1else:x=self.conv_false(x)x=tf.sigmoid(x)x=self.conv_shared(x)x=x-1x=self.conv_shared(x)returnx@nobuco.converter(ControlIf,channel_ordering_strategy=ChannelOrderingStrategy.FORCE_TENSORFLOW_ORDER)defconverter_ControlIf(self,x):order=ChannelOrder.TENSORFLOWkwargs={'inputs_channel_order':order,'outputs_channel_order':order,'return_outputs_pt':True}conv_pre,out_pre=nobuco.pytorch_to_keras(self.conv_pre,[x],**kwargs)conv_true,out_true=nobuco.pytorch_to_keras(self.conv_true,[out_pre],**kwargs)conv_false,out_false=nobuco.pytorch_to_keras(self.conv_false,[out_pre],**kwargs)conv_shared,_=nobuco.pytorch_to_keras(self.conv_shared,[out_true],**kwargs)layer=ControlIfKeras(conv_pre,conv_true,conv_false,conv_shared)returnlayerSeeexamplesfor other ways to convert control flow ops.Dynamic shapesWhat if we wanted our module to accept images of arbitrary height and width?
Can we have that? Let's try:classDynamicShape(nn.Module):def__init__(self):super().__init__()self.conv=nn.Conv2d(3,16,kernel_size=(1,1))defforward(self,x):x=self.conv(x)# Produces static shapeb,c,h,w=x.shapex=x[:,:,h//3:,w//3:]returnxinput=torch.normal(0,1,size=(1,3,128,128))pytorch_module=DynamicShape().eval()keras_model=nobuco.pytorch_to_keras(pytorch_module,args=[input],input_shapes={input:(None,3,None,None)},# Annotate dynamic axes with Noneinputs_channel_order=ChannelOrder.TENSORFLOW,outputs_channel_order=ChannelOrder.TENSORFLOW,)Something's not right. We don't see shape extraction ops in the debug output or the graph:That's not surprising, actually.
In Pytorch, tensor shape is a tuple of regular integers, not tensors, so it's quite difficult to track them.nobuco.shapesolves this problem.
This function returns tensors, much liketf.shapedoes:# Allows for dynamic shapeb,c,h,w=nobuco.shape(x)It's also possible to automatically substitute every.shape/.sizecall withnobuco.shapeduring the tracing phase by settingtrace_shapeflag:keras_model=nobuco.pytorch_to_keras(# ...trace_shape=True)Ad hoc modificationsLet's say, for illustrative purposes, that we prefer putting batchnormbeforeconvolution.
We were sureTFLiteConverterwould fuse these two linear operations into one.
Alas, it failed to meet our expectations.
Can we still get the fusion to work without re-training or messing around with the model checkpoint?classFusibleModule(nn.Module):def__init__(self):super().__init__()self.bn=nn.BatchNorm2d(3)self.conv=nn.Conv2d(3,16,kernel_size=(3,3),padding=(0,0))self.act=nn.ReLU()defforward(self,x):x=self.bn(x)x=self.conv(x)x=self.act(x)returnxHere's one way to do it:Wrap the two ops in aCallable. Decorate it with@nobuco.traceable.Make a custom converter for it which does the desired optimization.classFusibleModule(nn.Module):# ...@nobuco.traceabledefbn_conv(self,x):x=self.bn(x)x=self.conv(x)returnxdefforward(self,x):x=self.bn_conv(x)x=self.act(x)returnx@nobuco.converter(FusibleModule.bn_conv,channel_ordering_strategy=ChannelOrderingStrategy.FORCE_TENSORFLOW_ORDER)defconverter_bn_conv(self,x):order=ChannelOrder.TENSORFLOWbn,out_bn=nobuco.pytorch_to_keras(self.bn,[x],inputs_channel_order=order,outputs_channel_order=order,return_outputs_pt=True)conv=nobuco.pytorch_to_keras(self.conv,[out_bn],inputs_channel_order=order,outputs_channel_order=order)gamma,beta,moving_mean,moving_variance=bn.get_weights()kernel,bias=conv.get_weights()eps=self.bn.eps'''y = gamma * (x - moving_mean) / sqrt(moving_variance + eps) + betaz = kernel * y + bias=>z = kernel_fused * x + bias_fused WHEREkernel_fused = kernel * gamma / sqrt(moving_variance + eps)bias_fused = -kernel_fused * moving_mean + kernel * beta + bias'''kernel_fused=kernel*(gamma/np.sqrt(moving_variance+eps))[None,None,:,None]bias_fused=(-kernel_fused*moving_mean[None,None,:,None]+kernel*beta[None,None,:,None]).sum(axis=(0,1,2)).flatten()+biasconv.set_weights([kernel_fused,bias_fused])returnlambdaself,x:conv(x)So we put a converter inside your converterAs we've learned, Nobuco gets confused when in-place operation is applied to a slice.
There's a way to fix that, but let's not do it now.
Instead, we'll use it as an excuse to explain the concept of nested converters.
So, for this module, conversion will give us incorrect result:classSliceReLU(nn.Module):defforward(self,x):# Gives incorrect result after conversiontorch.relu_(x[:,1:2,16:25,8::2])# That's the recommended approach, but we're not going for it now# x[:, 1:2, 16:25, 8::2] = torch.relu_(x[:, 1:2, 16:25, 8::2])returnxclassMyModule(nn.Module):def__init__(self):super().__init__()self.conv=nn.Conv2d(3,3,kernel_size=(3,3),padding=(1,1))defforward(self,x):x=self.conv(x)SliceReLU()(x)returnxWe've seen it's possible to invoke a Nobuco converter inside another Nobuco converter.
Can we embed some third-party converter? You bet! Why? Because it might just do what we need.
Let's consider the standard route: Pytorch -> ONNX -> Tensorflow, with the latter step done withonnx-tf.
This library likes transposing stuff so much, converting the whole graph with it may introduce intolerable inference overhead. Nonetheless, it does the job.
A sensible tradeoff would be to wrap the problematic operation into its ownnn.Moduleand give it a special treat, while handling everything else with Nobuco.importonnxfromonnx_tf.backendimportprepare@nobuco.converter(SliceReLU,channel_ordering_strategy=ChannelOrderingStrategy.FORCE_PYTORCH_ORDER,reusable=False)defconverter_SliceReLU(self,x):model_path='slice_relu'onnx_path=model_path+'.onnx'# NB: onnx.export in implemented via tracing i.e. it may modify the inputs!torch.onnx.export(self,(x,),onnx_path,opset_version=12,input_names=['input'],dynamic_axes={'input':[0,1,2,3]})onnx_model=onnx.load(onnx_path)tf_rep=prepare(onnx_model)tf_rep.export_graph(model_path)model=tf.keras.models.load_model(model_path)returnkeras.layers.Lambda(lambdax:model(input=x))But was it worth it?Let's cut to the chase, here's the numbers.mobilenet_v3_large(26.8 Mb)nobucoonnx_tfspeedupx86 (XNNPACK)11.1 ms14.7 ms1.3xArm CPU (XNNPACK)24.3 ms40.3 ms1.6xArm GPU (OpenCL)21.3 ms192.6 ms9xdeeplabv3_resnet50(158.5 Mb)nobucoonnx_tfspeedupx86 (XNNPACK)1.25 s1.34 s1.07xArm CPU (XNNPACK)2.0 s2.7 s1.35xArm GPU (OpenCL)1.6 s2.6 s1.62xAs we can see, redundant transpositions may completely ruin the performance, especially on a GPU.
But that's not the only issue.
Let's test this:classSliceReLU(nn.Module):defforward(self,x):x[:,1:2,16:25,8::2]=torch.relu_(x[:,1:2,16:25,8::2])returnxnobucoonnx_tfspeedupx86 (XNNPACK)0.40 ms1.57 ms3.9xArm CPU4.6 ms2.9ms0.6xArm CPU (XNNPACK)2.1msFAIL—Arm GPU (OpenCL)21.8 msFAIL—Again, the graph obtained withonnx_tfis much slower on x86 CPU.
Worse yet, on mobile processor, optimized TFLite delegates for both GPU and CPU failed.
No transpose ops were added this time, so who's to blame?
Just look whattorch.onnx.exportgives us:slice_relu.onnxonnx_tfdoes a fair job optimizing the monstrosity of a graph it's given,
but combining consecutivesliceops seems to be too much to ask.
It also leaves out garbage nodes sometimes (note the free-floatingWhilein this example).Nobuco evades these types of problems by simply not dealing withonnx.slice_relu_nobucoslice_relu_onnxtfNobuco knowledge baseDon't want to convert anything but looking for a Tensorflow equivalent of a certain Pytorch node (operation or module)?
Nobuco already implements quite a few node converters, most written in concise and, hopefully, understandable way.
These are located innobuco/node_converters,
and there's a utility function to help you find what you need:node=torch.Tensor.repeat# node = F.relu_# node = nn.LSTMlocation_link,source_code=nobuco.locate_converter(node)print('Converter location:')print(location_link)print('Converter source code:')print(source_code)Converter location:File "/home/user/anaconda3/envs/nb/lib/python3.9/site-packages/nobuco/node_converters/tensor_manipulation.py", line 141Converter source code:@converter(torch.Tensor.repeat, channel_ordering_strategy=ChannelOrderingStrategy.MINIMUM_TRANSPOSITIONS)def converter_repeat(self, *sizes):def func(self, *sizes):if get_channel_order(self) == ChannelOrder.TENSORFLOW:sizes = permute_pytorch2keras(sizes)return tf.tile(self, sizes)return funcAcknowledgementsSlice assign converter is based onZaccharie Ramzi's tf-slice-assign script.