Torch模型转TensorFlow

Torch模型转TensorFlow

笔者最近在学习NLP方面的知识，发现很多论文的代码是用Torch实现的，Torch的模型部署没有TensorFlow成熟，笔者初步尝试Torch模型转Tensorflow模型，并采用TensorFlow Serving加载模型。笔者只是简单的学习实现了一个demo，后面尝试将其他的Torch模型转TensorFlow碰到了困难，没找到解决办法。下面介绍Torch模型转TensorFlow的模型的demo，主要包含以下内容：
（1）生成Torch模型
（2）Torch模型转onnx
（3）onnx转pb模型
（4）加载pb模型
（5）生成Tensorflow Serving加载的pb模型
（6）HTTP 请求
笔者的环境为torch==1.6.0，tensorflow==1.14.0，onnx==1.6.0，onnx-tf==1.5.0。

1. 生成Torch模型

这里写一个简单的分类模型，只有一层全连接层加一个sigmoid 输出层。代码如下：

import torch.nn as nn


class SimpleModel(nn.Module):
    def __init__(self, input_size, hidden_sizes, output_size):
        super(SimpleModel, self).__init__()
        self.input_size = input_size
        self.output_size = output_size
        self.fcs = []  
        in_size = input_size

        for i, next_size in enumerate(hidden_sizes):
            fc = nn.Linear(in_features=in_size, out_features=next_size)
            in_size = next_size
            self.__setattr__('fc{}'.format(i), fc)  
            self.fcs.append(fc)

        self.last_fc = nn.Linear(in_features=in_size, out_features=output_size)

    def forward(self, x):
        for i, fc in enumerate(self.fcs):
            x = fc(x)
            x = nn.ReLU()(x)
        out = self.last_fc(x)
        return nn.Sigmoid()(out)

模型训练如下，训练集和测试集的数据都是随机数。

model_pytorch = SimpleModel(input_size=input_size, hidden_sizes=hidden_sizes,
                            output_size=output_size)
model_pytorch = model_pytorch.to(device)

# Set loss and optimizer
# Set binary cross entropy loss since 2 classes only
criterion = nn.BCELoss()
optimizer = optim.Adam(model_pytorch.parameters(), lr=1e-3)

num_epochs = 30

# Train model
time_start = time.time()

for epoch in range(num_epochs):
    model_pytorch.train()

    train_loss_total = 0

    for data, target in train_loader:
        data, target = data.to(device), target.float().to(device)
        optimizer.zero_grad()
        output = model_pytorch(data)
        train_loss = criterion(output, target)
        train_loss.backward()
        optimizer.step()
        train_loss_total += train_loss.item() * data.size(0)

    print('Epoch {} completed. Train loss is {:.3f}'.format(epoch + 1, train_loss_total / train_size))
print('Time taken to completed {} epochs: {:.2f} minutes'.format(num_epochs, (time.time() - time_start) / 60))

# Evaluate model
model_pytorch.eval()

test_loss_total = 0
total_num_corrects = 0
threshold = 0.5
time_start = time.time()

for data, target in test_loader:
    data, target = data.to(device), target.float().to(device)
    optimizer.zero_grad()
    output = model_pytorch(data)
    train_loss = criterion(output, target)
    train_loss.backward()
    optimizer.step()
    train_loss_total += train_loss.item() * data.size(0)

    pred = (output >= threshold).view_as(target)  
    num_correct = torch.sum(pred == target.byte()).item()
    total_num_corrects += num_correct

print('Evaluation completed. Test loss is {:.3f}'.format(test_loss_total / test_size))
print('Test accuracy is {:.3f}'.format(total_num_corrects / test_size))
print('Time taken to complete evaluation: {:.2f} minutes'.format((time.time() - time_start) / 60))

if not os.path.exists('./models/'):
    os.mkdir('./models/')

torch.save(model_pytorch.state_dict(), './models/model_simple.pt')

2. Torch模型转onnx

将上述模型训练保存的模型转为onnx，代码如下：

model_pytorch = SimpleModel(input_size=input_size,
                            hidden_sizes=hidden_sizes,
                            output_size=output_size)
model_pytorch.load_state_dict(torch.load('./models/model_simple.pt'))
dummy_input = torch.randn([1, 20])
dummy_output = model_pytorch(dummy_input)
print(dummy_output)

# Export to ONNX format
torch.onnx.export(model_pytorch, dummy_input, './models/model_simple.onnx',
                  input_names=['test_input'],
                  output_names=['test_output'])

备注：如果转onnx报错，根据实际情况修改torch.onnx.export()中的参数

3. onnx转pb模型

利用onnx将Torch模型转为TensorFlow模型。代码如下：

import onnx
from onnx_tf.backend import prepare

model_onnx = onnx.load('./models/model_simple.onnx')
tf_rep = prepare(model_onnx)
# Export model as .pb file
tf_rep.export_graph('./models/model_simple.pb')

备注：onnx得到的pb模型，不能用Tensorflow Serving加载，需要进一步将p模型转为Tensorflow Serving加载的pb模型。

4. 加载pb模型

加载pb模型做预测，代码如下：

import numpy as np
import tensorflow as tf


def load_pb(path_to_pb):
    with tf.gfile.GFile(path_to_pb, 'rb') as f:
        graph_def = tf.GraphDef()
        graph_def.ParseFromString(f.read())
    with tf.Graph().as_default() as graph:
        tf.import_graph_def(graph_def, name='')
        return graph


tf_graph = load_pb('./models/model_simple.pb')
sess = tf.Session(graph=tf_graph)

# Show tensor names in graph
for op in tf_graph.get_operations():
  print(op.values())

output_tensor = tf_graph.get_tensor_by_name('test_output:0')
input_tensor = tf_graph.get_tensor_by_name('test_input:0')

dummy_input = np.random.randn(1, 20).astype(np.float32)
output = sess.run(output_tensor, feed_dict={input_tensor: dummy_input})

5. 生成Tensorflow Serving加载的pb模型

def generator_tf_serving_pb_v1(export_dir, graph_pb):
    builder = tf.saved_model.builder.SavedModelBuilder(export_dir)
    with tf.gfile.GFile(graph_pb, "rb") as f:
        graph_def = tf.GraphDef()
    graph_def.ParseFromString(f.read())


    sigs = {}
    with tf.Session(graph=tf.Graph()) as sess:
        tf.import_graph_def(graph_def, name="")

        g = tf.get_default_graph()

        sigs[signature_constants.DEFAULT_SERVING_SIGNATURE_DEF_KEY] = \
            tf.saved_model.signature_def_utils.predict_signature_def(
                inputs={"test_input": g.get_tensor_by_name('test_input:0')},
                outputs={"output": g.get_tensor_by_name('test_output:0')}
            )

        builder.add_meta_graph_and_variables(sess, [tag_constants.SERVING],
                                             signature_def_map=sigs)

        builder.save()

6. HTTP请求

将TensorFlow Serving加载5.生成Tensorflow Serving加载的pb模型得到的pb模型，采用HTTP服务请求的方式得到模型的预测结果，代码如下：

@app.route('/mlp', methods=['POST'])
def mlp_model_infer():
    params = json.loads(request.get_data(), encoding="utf-8")
    tf_url = params.get("url")
    input_ = np.random.randn(1, 20).astype(np.float32).tolist()

    tensor = {"instances": [{"test_input": input_}]}
    result = requests.post(tf_url, json=tensor)
    if result.status_code == 200:
        pred = result.json()['predictions'][0]
        reture_result = {"code": 200,
                         "message": "finish",
                         "result": pred}
        return jsonify(reture_result)
    else:
        reture_result = {"code": 200,
                         "message": "faile"}
        return jsonify(reture_result)

以上是笔者的Torch 模型转pb模型的demo，如果有误，欢迎大家指出。笔者尝试将其他的Torch模型转pb失败，还在解决问题中。