Pytorch和tensorflow api映射

0 前言

当前pytroch是最流行的深度学习框架， 最主要的优点就是灵活易用。 tensorflow在很多生产部署环境中应用也还比较广泛， 虽然不好用， 但是效率上还是有点优势的。很多最新的研究工作是基于torch做的， 实际应用时可能需要转换为tensorflow。 因此， 这里列出一些常用的api映射关系，方便查阅和避坑。

需要注意的是， torch的api比较简单， 但tensorflow的api却非常的混乱， 有底层api， 高级api， 还有各种第3方库的api， 功能上很多是重复的。并且tf 1.x版本与tf 2.x版本也有非常大的区别。 为了简单起见， 建议尽量使用tf.keras中的高级api， 这些api跟torch也比较近似。 这里也只考虑tf 1.x版本。

torch api : https://pytorch-cn.readthedocs.io/zh/latest/package_references/torch-nn/
tf api: https://www.tensorflow.org/versions/r1.15/api_docs/python/tf?hl=zh-cn

1 torch.nn

torch	tf	不同点
nn.Conv2d	tf.keras.layers.Conv2D	tf接口没有输入通道，group这些参数， padding是选择类型，不是padding数，并且在stride>1时， 以及padding数左右不对称的时候，padding的策略是不一样的（详见注2）, 另外data_format默认是’channels_last’, 跟torch不一致
nn.BatchNorm2d	tf.keras.layers.BatchNormalization
nn.PReLU	tf.keras.layers.PReLU	控制参数不同
nn.ReLU	tf.keras.ReLU	tf没有inplace参数
nn.Sequential	tf.keras.Sequential	tf要求输入是list
nn.Sigmoid	tf.keras.activations.sigmoid	tf的api是个函数
nn.AvgPool2d	tf.keras.layers.AveragePooling2D
nn.Softmax	tf.keras.layers.Softmax
nn.Upsample	tf.keras.layers.UpSampling2D
nn.Pixshuffle	tf.depth_to_space	tf的api是个函数， 注意： 这里还有一个功能等价上的大坑， 详见 tf depth_to_space 与torch pixshuffle 踩坑

2 函数

torch	tf	不同点
torch.cat	tf.concat	tf的参数是axis, torch 是dim
F.grid_sample	可参考https://www.zhihu.com/question/432735000 实现， 实测发现第一个的输出对不上， 第二个可以对上	第一个输入格式不同， torch是NCHW, tf是NHWC
F.interpolate	tf.image.resize_bilinear tf.image.resize_nearest_neighbor	F.interpolate有很多插值方法，tf中有多个函数与之对应.需要特别注意的是， tf中的操作是按NHWC的数据排布操作的， 如果输入是NCHW的， 要特别注意需要转一下， 不然会在C和H维上进行插值; 另外有一个地方需要特别注意，tf中有两个控制参数， align_corners和half_pixel_centers , torch只有align_corners。 torch的align_corners=False与tf的align_corners=False，half_pixel_centers=True 是等价的， torch的align_corners=True与tf的align_corners=True，half_pixel_centers=False 是等价的
unsqueeze	tf.expand_dims	torch中该方法是Tensor的属性， 但tf中只是个函数, torch参数名是dim,tf是perm
transpose	tf.transpose	torch中该方法是Tensor的属性， 但tf中只是个函数, torch参数名是dim, tf是axis
torch.arange	tf.keras.backend.arange
torch.meshgrid	tf.meshgrid	需要特别注意， tf中有indexing这个参数， 且默认是xy,做了转置，indexing为ij时才和torch一致
torch.stack	tf.stack	tf的参数是axis, torch 是dim
F.pad	tf.pad	参数名pad不一致，torch是pad, tf是paddings, 尤其要注意的是pad的含义也不同; mode的取值不同， 详见注[1]
shape	tf.shape
view	tf.reshape
var mean	tf.nn.moments	tf可以同时求均值和方差，参数是axes
sqrt	tf.math.sqrt
expand	tf.tile tf.repeat	没有直接等价的算子， 但tile 和repeat可以做到expand的功能， tile使用更简单一些

注:
[1] torch中pad中是一个tuple， 元素个数是2的整数倍， 最多为Nx2, 其中N是输入Tensor的rank， 如NCHW的输入， rank是4。 pad的含义是， 维度从后往前， 每一个维度有2个pad值(左和右)。
举例来说，对于一个NCHW的输入， pad=（1,1）表示在W维度上左右各pad一个值， pad=(1,1,2,2)表示在W维上左右各pad 一个值， 在H维上左右各pad两个值。
tf中paddding是一个N*2的tensor， 如NCHW的输入， padding就是一个4x2的Tensor. padding的含义就非常明确，每一维2个padding值。如对于NCHW输入padding=[[0,0],[0,0],[2,2],[1,1]]表示在H维度上左右各padding两个值， 在W维度左右各padding1个值。
因为通常都只在HW维度上做padding， 所以torch写的会更少一点， 但要注意维度是从后往前的。

[2] tf的卷积与torch的不同
tf中没有group卷积， 只有depthwise conv这个特殊的分组卷积，实现的话， 可以自行分组， 用for循环的方式， 每次处理一组卷积， 然后把结果拼接起来。

tf中conv的padding方式与torch是不同的。 首先， tf中的padding只能选择valid 和same两种模式， 不像torch可以灵活进行padding。 不过一般用的多也就是valid和same。 另外当padding左右不对称时， padding的策略是不一样的， 比如kernel_size=2, stride=1, 要想输出大小保持不变， 只需要padding 1行(列)值。那么问题来了， padding在哪个位置呢？ torch中是padding在左边和上边， 而tf是padding在右边和下边。 这就导致输出结果无法对齐， 解决办法就是tf中用tf.pad自行进行padding， 然后定义卷积时padding=valid。

对stride>1的情况， padding也有不同。 比如， 我在实际中遇到这样一组参数， kernel_size=7, stride=4, 输出保持大小不变。（其实不变是相对stride=1的， stride>1时尺寸肯定是变小的， 为W/stride)。按理说这时需要padding的个数是6， 是偶数， 应该是左右上下各padding3个， 但是实际中tf卷积结果和torch差距巨大。 采用上面的方法， 手动在左右上下各padding 3个， 结果就对上了。
所以不是特别清楚tf里面padding的规则。 如果有了解比较清楚的小伙伴， 欢迎留言沟通。