tensorflow for python做模型训练、tensorflow for java做模型预测(同时生成pb文件和variable变量)

亦凉 2021-09-28 09:10 342阅读 0赞

python脚本(此代码为线性回归的demo)

\#!/usr/bin/python  
\# -\*- coding:utf-8 -\*-

import tensorflow as tf  
from tensorflow import saved\_model as sm  
import numpy as np

x\_data = np.linspace(-0.5,0.5,200)\[:,np.newaxis\]  
noise = np.random.normal(0,0.02,x\_data.shape)  
y\_data = np.square(x\_data) + noise

x = tf.placeholder(tf.float32,\[None,1\])  
y = tf.placeholder(tf.float32,\[None,1\])

Weights\_L1 = tf.Variable(tf.random\_normal(\[1,10\]))  
biases\_L1 = tf.Variable(tf.zeros(\[1,10\]))  
Wx\_plus\_b\_L1 = tf.matmul(x,Weights\_L1) + biases\_L1  
L1 = tf.nn.tanh(Wx\_plus\_b\_L1)

Weights\_L2 = tf.Variable(tf.random\_normal(\[10,1\]))  
biases\_L2 = tf.Variable(tf.zeros(\[1,1\]))  
Wx\_plus\_b\_L2 = tf.matmul(L1,Weights\_L2) + biases\_L2  
prediction = tf.nn.tanh(Wx\_plus\_b\_L2)

loss = tf.reduce\_mean(tf.square(y - prediction))  
\# 使用梯度下降法训练  
train\_step = tf.train.GradientDescentOptimizer(0.1).minimize(loss)

with tf.Session() as sess:  
    \# 变量初始化  
    sess.run(tf.global\_variables\_initializer())  
    for \_ in range(2000):  
        sess.run(train\_step, feed\_dict=\{x: x\_data, y: y\_data\})  
          
    \# 获得预测值  
    prediction\_value = sess.run(prediction, feed\_dict=\{x: x\_data\})  
    print(prediction\_value)

\# 1、构建SavedModelBuilder的实例builder，并设置模型导出路径  
    path = 'hdfs://default/home/hdp\_ubu\_tech\_wei/resultdata/wuxian\_prs/hometown/hezhidong/LinearModel'  
    builder = sm.builder.SavedModelBuilder(path)

\# 2、定义模型服务的SignatureDef protobuf  
    X\_TensorInfo = sm.utils.build\_tensor\_info(x)  
    y\_TensorInfo = sm.utils.build\_tensor\_info(prediction)  
    SignatureDef = sm.signature\_def\_utils.build\_signature\_def(  
        inputs=\{'input\_1': X\_TensorInfo\},  
        outputs=\{'output': y\_TensorInfo\},  
        method\_name=sm.signature\_constants.REGRESS\_METHOD\_NAME  
    )

\# 3、将 graph 和变量等信息写入 MetaGraphDef protobuf  
    builder.add\_meta\_graph\_and\_variables(sess, tags=\[tf.saved\_model.tag\_constants.SERVING\],  
                                         signature\_def\_map=\{'prediction':SignatureDef,\}  
                                         )  
    builder.save()