pytorch model to keras model

概述

使用pytorch建立的模型，有时想把pytorch建立好的模型装换为keras，本人使用TensorFlow作为keras的backend

依赖

标准库依赖：

  1. pytorch
  2. keras
  3. tensorflow
  4. pytorch2keras
安装方式
conda install tensorflow-gpu  keras
conda install pytorch torchvision
pip install pytorch2keras

代码

去pytorch2keras页面，具体的 代码片.

import numpy as np
import torch
from torch.autograd import Variable
from pytorch2keras import converter

class Pytorch2KerasTestNet(torch.nn.Module):
    def __init__(self):
        super(Pytorch2KerasTestNet, self).__init__()
        self.conv1 = ConvLayer(3, 32, kernel_size=9, stride=1)
        self.in1 = torch.nn.InstanceNorm2d(32, affine=True)
        self.relu = torch.nn.ReLU()

    def forward(self, x):
        y = self.relu(self.in1(self.conv1(x)))
        return y


class ConvLayer(torch.nn.Module):
    def __init__(self, in_channels, out_channels, kernel_size, stride):
        super(ConvLayer, self).__init__()
        reflection_padding = kernel_size // 2
        self.reflection_pad = torch.nn.ReflectionPad2d(reflection_padding)
        self.conv2d = torch.nn.Conv2d(in_channels, out_channels, kernel_size, stride)

    def forward(self, x):
        out = self.reflection_pad(x)
        
        print("conv2d")
        out = self.conv2d(out)
        return out
        
model   = Pytorch2KerasTestNet()

input_np = np.random.uniform(0, 1, (1, 3, 224, 224))
input_var = Variable(torch.FloatTensor(input_np))
k_model = converter.pytorch_to_keras(model, input_var, [(3, 224, 224,)], verbose=True) 
k_model.summary() 

#保存模型
k_model.save('my_model.h5')

# 重新载入模型
from keras.models import load_model
import tensorflow as tf

model = load_model('my_model.h5',custom_objects={"tf": tf})
model.summary()

输出结果
pytorch Model to keras model
[1]: pytorch2keras 参考
[2]: pytorch 安装
[3]: keras 载入模型出错 custom_objects