pytorch模型迁移到tensorflow

高晋峰

2019年09月21日 00:00

pytorch / tensorflow

前几天用pytorch训练了一个目标检测的模型，但是因为要部署的嵌入式板子只支持tensorlfow、caffe等模型，最重要的是我从来没有使用过这两框架，坑跌的是tensorflow也只有1.9版本支持。想到了四种方案，第一种tensorlfow重新训练编写模型、第二种使用onnx框架进行模型转换、第三种使用caffe重写整个模型、第四使用板子支持的框架只重写inference阶段的模型，同时将pytorch训练好的权重进行迁移

想到这四种方案为了偷懒肯定是直接用onnx最简单啦直接pytorch导出onnx然后转换成tensoflow立马此时一下ojbk啊，赶紧部署，我日不支持onnx。。。。退而求其次听说caffe挺好用的，一看安装要从源码编译我耸了，看看tensorflow吧，这个挺火的一看官网这么多版本而且文档。。。我曹莫得办法只能硬着头皮学tensorflow了，看了半天莫烦的视频初步对tf的静态图结构、编写程序的流程有了一个大概的了解，直接区官网copy了一个教程，对着pytorch中定义的网络结构照猫画虎写了tf版的模型，一个大概二十层的轻量级检测网络打起来了，最终的目的就是使用tensorflow重写了整个网络结构。下一步就是将torch中训练好的参数迁移过来，此过程断断续续经历了好几天，途中各种坑啊，再次简单记录一下

第一坑：就是图片在内存中存储的格式，pytorch默认NCHW，tf默认NHWC
第二坑：卷积padding的坑
第三坑：pyotrch参数 convert 成tf参数的坑
第四坑：batch normalization的坑（真坑！！）

第一坑：NCHW VS NHWC

需要在tensorflow实现卷积的过程中手动设置卷积conv2d参数以及bias参数的data_format = “NCHW”，相应的需要将stride中每一维度的值与数据格式相对应，比如说HW方向stride为2，如果data_format为NHWC则 stride = [1, 2, 2, 1]，修改数据存储格式为NCHW之后相应的 stride = [1, 1, 2, 2]。还有就是卷积核的shape。 tf中卷集核是按照(K_H, K_W, CH_In, CH_out)的形状组织的，但是pytorch中Kernel的shape是( CH_In, CH_out, K_H, K_W )。因此在将torch中的卷积核赋值到 tf卷积核的时候应该注意首先将torch中的Kernel进行transpose：具体如下代码：

1	tf_kernel = np.transpose(pytorch_kernel, (2, 3, 1, 0))

还有就是对于BN层的参数torch中的 xxx.norm.weight”, xxx.norm.bias”对应的tf参数中的xxxx/batch_normalization/beta”, xxxx/batch_normalization/gamma”。

第二坑：padding的区别

tf中padding有两种VALID和SAME

对于VALID，就是input不做任何的padding可能会丢弃输入图片最后的某几行和列，对应的pytorch中就是设置卷积参数 padding=0

对于SAME，tf中SAME的意思就是在 stride = 1, 1的时候输入输出的大小下面的公式计算得出：

1
2
3

For the `SAME` padding, the output height and width are computed as:  
out_height = ceil(float(in_height)  / float(strides[1]))
out_width  = ceil(float(in_width)   / float(strides[2]))

补零的方式不同


`tensorflow ："SAME"` = with zero padding:

                   pad|                                      |pad
       inputs:      0 |1  2  3  4  5  6  7  8  9  10 11 12 13|0  0
                   |________________|
                                  |_________________|
                                                 |________________|

    Input width = 13 Filter width = 6 Stride = 5

我们计算这个example中 out_height = ceil(float(in_height) / float(strides[1])) 带入值，out_height = ceil（13 / 5） = 3. 此时通过计算需要padding的大小，要使得（x - 6）/ 5 + 1 = 3, 得出x = 16,需要padding的大小为3， tensorflow首先会尝试平均的在输入的左右进行0 padding，但如果需要padding的值为奇数，则在输入的右边padding，上面的例子中可以看到padding是3左右各padding一个0之后剩余的一个0加到了右边，也就是说右边的优先级高。对于上下方向的padding，同理可得优先padding下方向。

而在pytorch中padding参数是一个长度为的tuple（h，w），分别代表着row方向和column方向同时padding的行数和列数。也就是说如果padding=[1 ,1]。pytorch将会在输入图片的左右上下同时padding 1 步，和tf不同的是torch中是成对的padding不存在只padding行右而没有左边，只上不下的情况。但是可以使用torch.nn.functional.pad函数手动设置(padLeft, padRight, padTop, padBottom)各个方向的padding步数。通过查看源代码并且自己动手做小实验比较结果得出的。

第三坑：参数转换

参数转换的过程中一定要注意各个参数的对应关系以及 shape一定要assert一下。刚开始的时候由于粗心在transpose的时候

1 2	tf_kernel = np.transpose(pytorch_kernel, (2, 3, 1, 0)) 错误的写成了 tf_kernel = np.transpose(pytorch_kernel, (3, 2, 1, 0))

结果最后的计算结果总是对不上，简直了太傻比了。还有就是参数名要对应其实也挺麻烦的，要么在写模型的时候注意与torch中的参数名保持一致，要么就是自己根据某着对应规则在参数迁移的时候保证一一对应，这里我用一个字典（key就是tf参数名，value就是对应的torch参数名），然后手动调整对应关系也算比较方便吧，但是当参数量很大的时候建议还是使用命名空间一致对应来解决。

第四坑：tf和pytorch中的bn参数设置大坑

1 2	# pytorch: class BatchNorm2d(num_features, eps=1e-5, momentum=0.1, affine=True, track_running_stats=True) # tf：def batch_normalization(inputs,axis=-1,momentum=0.99,epsilon=1e-3, ...)

刚开始没仔细看tf中bn的参数默认值后来调试发现bn计算结果总是对不上，将tf中epsilion改为torch中对应的1e-5, momentum修改要特别注意！！！tf中momentum 其实对应的是torch中的 1 - momentum， torch中原本是0.1,所以将tf中momentum修改为0.9即可.

总结一下

踩坑过程中一定要细心调试，多看源代码做实验验证自己的想法期间为了验证脑子里稀奇古怪的想法用numpy手写了一个conv2d来验证各个框架之间的conv2d方式。此外特别要注意的是调试一定要有耐心细心，有时结果可能会有非常细微的差别，一定要做好对比结果！！比如刚开始把bn的beta和gamma参数弄反了，发现计算结果很相近但是小数点后几位确差的很多，很奇怪后来检查参数发现beta和gamma有一组搞反了。

注意

torch.version = ‘1.0.1’ ， tensorflow.version = ‘1.9.0’ 均为gpu版本。

参考文献

https://stackoverflow.com/questions/37674306/what-is-the-difference-between-same-and-valid-padding-in-tf-nn-max-pool-of-t

YNZH's Blog