PyTorch到Tensorflow的模型转换
文章目录
- 1 结果相同，使用ONNX的不同框架
- 2 PyTorch模型转换pipeline
- 3 转移模型结果
在本文中，我们将学习如何将PyTorch模型转换为TensorFlow。
如果您是深度学习的新手，那么使用哪种框架可能会让您不知所措。我们个人认为PyTorch是您应该学习的第一个框架，但它可能不是您想要学习的唯一框架。
好消息是您不需要嫁给框架。您可以在PyTorch中训练模型，然后将其轻松转换为Tensorflow，只要使用标准图层即可。实现此转换的最佳方法是先将PyTorch模型转换为ONNX，然后再转换为Tensorflow/Keras格式。
1 结果相同，使用ONNX的不同框架
我们可以观察到，在有关FCN ResNet-18 PyTorch的早期文章中，所实现的模型比TensorFlow FCN版本更准确地预测了图片中的单峰骆驼区域：
PyTorch FCN ResNet18激活：
TensorFlow FCN ResNet50激活：
假设我们想捕获结果并将它们转移到另一个领域，例如，从PyTorch到TensorFlow。有什么办法可以执行吗？答案是肯定的。一种可能的方法是使用pytorch2keras库。如其名称所述，该工具提供了在PyTorch和Keras等框架之间进行模型转换的简便方法。您可以使用pip轻松安装它：
pip3 install pytorch2keras
2 PyTorch模型转换pipeline
从pytorch2keras的仓库中可以看到，converter.py中描述了 pipeline 的逻辑。让我们查看其关键点：
def pytorch_to_keras(
model, args, input_shapes=None,
change_ordering=False, verbose=False, name_policy=None,
use_optimizer=False, do_constant_folding=False
):
# ...
# load a ModelProto structure with ONNX
onnx_model = onnx.load(stream)
# ...
#
k_model = onnx_to_keras(onnx_model=onnx_model, input_names=input_names,
input_shapes=input_shapes, name_policy=name_policy,
verbose=verbose, change_ordering=change_ordering)
return k_model
您可能已经注意到，该工具基于Open Neural Network Exchange (ONNX)。ONNX是一个开源AI项目，其目标是使不同工具之间的神经网络模型互换，以选择这些工具的更好组合。将获得的过渡的top-level ONNX ModelProto容器传递给onnx2keras工具的onnx_to_keras函数，以进行进一步的层映射。
让我们以全卷积网络架构为例来研究PyTorch ResNet18转换过程：
# import transferring tool
from pytorch2keras.converter import pytorch_to_keras
def converted_fully_convolutional_resnet18(
input_tensor, pretrained_resnet=True,
):
# define input tensor
input_var = Variable(torch.FloatTensor(input_tensor))
# get PyTorch ResNet18 model
model_to_transfer = FullyConvolutionalResnet18(pretrained=pretrained_resnet)
model_to_transfer.eval()
# convert PyTorch model to Keras
model = pytorch_to_keras(
model_to_transfer,
input_var,
[input_var.shape[-3:]],
change_ordering=True,
verbose=False,
name_policy="keep",
)
return model
现在我们可以比较PyTorch和TensorFlow FCN版本了。让我们看一下第一堆PyTorch FullyConvolutionalResnet18层。值得注意的是，我们使用了torchsummary工具来实现PyTorch和TensorFlow模型摘要的视觉一致性：
from torchsummary import summary
summary(model_to_transfer, input_size=input_var.shape[-3:])
输出为：
----------------------------------------------------------------
Layer (type) Output Shape Param #
================================================================
Conv2d-1 [-1, 64, 363, 960] 9,408
BatchNorm2d-2 [-1, 64, 363, 960] 128
ReLU-3 [-1, 64, 363, 960] 0
MaxPool2d-4 [-1, 64, 182, 480] 0
Conv2d-5 [-1, 64, 182, 480] 36,864
BatchNorm2d-6 [-1, 64, 182, 480] 128
ReLU-7 [-1, 64, 182, 480] 0
Conv2d-8 [-1, 64, 182, 480] 36,864
BatchNorm2d-9 [-1, 64, 182, 480] 128
ReLU-10 [-1, 64, 182, 480] 0
BasicBlock-11 [-1, 64, 182, 480] 0
Conv2d-12 [-1, 64, 182, 480] 36,864
BatchNorm2d-13 [-1, 64, 182, 480] 128
ReLU-14 [-1, 64, 182, 480] 0
Conv2d-15 [-1, 64, 182, 480] 36,864
BatchNorm2d-16 [-1, 64, 182, 480] 128
ReLU-17 [-1, 64, 182, 480] 0
使用pytorch_to_keras函数转换后获得的TensorFlow模型包含与初始PyTorch ResNet18模型相同的层，但TF特定的InputLayer和ZeroPadding2D作为填充参数包含在torch.nn.Conv2d中。
以下摘要是使用tf.keras.Model类的内置Keras summary方法生成的：
model.summary()
输出中的相应层已标记有用于PyTorch-TF映射的适当数字：
Layer (type) Output Shape Param #
===============================================================
input_0 (InputLayer) [(None, 725, 1920, 3 0
_______________________________________________________________
125_pad (ZeroPadding2D) (None, 731, 1926, 3) 0
_______________________________________________________________
125 (Conv2D) (None, 363, 960, 64) 9408 1
_______________________________________________________________
126 (BatchNormalization) (None, 363, 960, 64) 256 2
_______________________________________________________________
127 (Activation) (None, 363, 960, 64) 0 3
_______________________________________________________________
128_pad (ZeroPadding2D) (None, 365, 962, 64) 0
_______________________________________________________________
128 (MaxPooling2D) (None, 182, 480, 64) 0 4
_______________________________________________________________
129_pad (ZeroPadding2D) (None, 184, 482, 64) 0
_______________________________________________________________
129 (Conv2D) (None, 182, 480, 64) 36864 5
_______________________________________________________________
130 (BatchNormalization) (None, 182, 480, 64) 256 6
_______________________________________________________________
131 (Activation) (None, 182, 480, 64) 0 7
_______________________________________________________________
132_pad (ZeroPadding2D) (None, 184, 482, 64) 0
_______________________________________________________________
132 (Conv2D) (None, 182, 480, 64) 36864 8
_______________________________________________________________
133 (BatchNormalization) (None, 182, 480, 64) 256 9
_______________________________________________________________
134 (Add) (None, 182, 480, 64) 0
_______________________________________________________________
135 (Activation) (None, 182, 480, 64) 0 10
_______________________________________________________________
136_pad (ZeroPadding2D) (None, 184, 482, 64) 0
_______________________________________________________________
136 (Conv2D) (None, 182, 480, 64) 36864 12
_______________________________________________________________
137 (BatchNormalization) (None, 182, 480, 64) 256 13
_______________________________________________________________
138 (Activation) (None, 182, 480, 64) 0 14
_______________________________________________________________
139_pad (ZeroPadding2D) (None, 184, 482, 64) 0
_______________________________________________________________
139 (Conv2D) (None, 182, 480, 64) 36864 15
_______________________________________________________________
140 (BatchNormalization) (None, 182, 480, 64) 256 16
_______________________________________________________________
141 (Add) (None, 182, 480, 64) 0
_______________________________________________________________
142 (Activation) (None, 182, 480, 64) 0 17
_______________________________________________________________
143_pad (ZeroPadding2D) (None, 184, 482, 64) 0
_______________________________________________________________
以下方案部分介绍了两种版本的TensorFlow和PyTorch的FCN ResNet18块的可视化表示：
左侧为转换后的TensorFlow FCN ResNet18模型，右侧为初始的PyTorchFCN ResNet18模型。
使用Netron开源查看器生成模型图。它支持从ONNX，TensorFlow，Caffe，PyTorch等获得的各种模型格式。保存的模型图作为输入传递给Netron，Netron进一步生成详细的模型图。
3 转移模型结果
因此，我们将整个PyTorch FC ResNet-18模型及其权重转换为TensorFlow，将NCHW（批次大小，通道，高度，宽度）格式更改为NHWC，并带有change_ordering = True参数。
这样做是因为在PyTorch模型中，输入层的 shape 为3×725×1920，而在TensorFlow中，其输入层的 shape 已更改为725×1920×3，因为TF中的默认数据格式为NHWC。我们还应该记住，要获得与PyTorch (1, 1000, 3, 8)相同的预测shape，我们应该再次转置网络输出：
# NHWC: (1, 725, 1920, 3)
predict_image = tf.expand_dims(image, 0)
# NCHW: (1, 3, 725, 1920)
image = np.transpose(tf.expand_dims(image, 0).numpy(), [0, 3, 1, 2])
# get transferred torch ResNet18 with pre-trained ImageNet weights
model = converted_fully_convolutional_resnet18(
input_tensor=image, pretrained_resnet=True,
)
# Perform inference.
# Instead of a 1×1000 vector, we will get a
# 1×1000×n×m output ( i.e. a probability map
# of size n × m for each 1000 class,
# where n and m depend on the size of the image).
preds = model.predict(predict_image)
# NHWC: (1, 3, 8, 1000) back to NCHW: (1, 1000, 3, 8)
preds = tf.transpose(preds, (0, 3, 1, 2))
preds = tf.nn.softmax(preds, axis=1)
要提到的另一点是图像预处理。我们记得在TF全卷积ResNet50中，应用了特殊的preprocess_input util函数。但是，在这里，为了转换为TF模型，我们使用与PyTorch FCN ResNet-18情况相同的归一化方法：
# transform input image:
transform = Compose(
[
Normalize(
# subtract mean
mean=(0.485, 0.456, 0.406),
# divide by standard deviation
std=(0.229, 0.224, 0.225),
),
],
)
# apply image transformations, (725, 1920, 3)
image = transform(image=image)["image"]
让我们来探讨结果：
Response map shape : (1, 1000, 3, 8)
Predicted Class : Arabian camel, dromedary, Camelus dromedarius tf.Tensor(354, shape=(), dtype=int64)
预测的类别是正确的，让我们看一下响应图：
您可以看到，响应区域与之前的PyTorch FCN帖子中的响应区域相同。
转换后的TF FCN ResNet18结果
