摘要:GraphSAGE
,tensorflow
,tensorflow_model_server
,tensorboard
,saved_model_cli
GraphSAGE实战目标
GraphSAGE的目标是对节点的邻居采样,从而避免每次计算都需要拿到全部节点的邻接矩阵,因此可以将训练好的模型直接用于新节点预测,GraphSAGE的实战目标就是对新的在训练数据中没有出现过的中心节点,基于该节点自身的特征和邻居特征,预测该节点的任务,比如预测该节点的分类。
GraphSAGE数据链路分析
这块参考另一篇博客文章https://www.jianshu.com/p/073e39c2d81a,简单而言就是模型层实际上就是学习到了每一层卷积的W,以及最后一层全连接的W,输入部分输入节点自身
,节点1跳
,节点2跳
的特征向量
,整体来看还是比较清晰简单的。
数据准备预处理
数据采用的cora数据,将全部数据分为训练
,验证
,预测
,三个部分独立存储,其中验证集作为早停,预测集用来评价模型,注意虽然三个数据集的节点索引互不重合,但是在采样过程中训练集的节点可以采样到他的属于验证集的邻居,只是拿到邻居的特征向量,不知道节点的y值。具体实现如下:
import os
import pickle
import numpy as np
import scipy.sparse as sp
BASIC_PATH = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
def data_split():
names = ['x', 'y', 'tx', 'ty', 'allx', 'ally', 'graph']
objects = []
for i in range(len(names)):
with open(os.path.join(BASIC_PATH, "./data/ind.cora.{}".format(names[i])), 'rb') as f:
objects.append(pickle.load(f, encoding='latin1'))
x, y, tx, ty, allx, ally, graph = tuple(objects)
test_idx_reorder = [int(x.strip()) for x in open(os.path.join(BASIC_PATH, "./data/ind.cora.test.index"), "r").readlines()]
test_idx_range = np.sort(test_idx_reorder)
# 测试索引位置修正
features = sp.vstack((allx, tx)).tolil()
features[test_idx_reorder, :] = features[test_idx_range, :]
labels = np.vstack((ally, ty))
labels[test_idx_reorder, :] = labels[test_idx_range, :]
# 训练[:1000],验证[1000:1708],测试[1708:]
train_nodes = list(range(1000))
train_y = labels[train_nodes]
val_nodes = list(range(1000, 1708))
val_y = labels[val_nodes]
test_nodes = list(range(1708, 2708))
test_y = labels[test_nodes]
return train_nodes, train_y, val_nodes, val_y, test_nodes, test_y, graph, features
def sample(nodes, neighbour_list, k=2, num_supports=None):
if num_supports is None:
num_supports = [10, 25]
assert len(num_supports) == k, "num_supports长度必须和k阶相等"
layer_neighbours = {}
for i in range(k):
neighbours = []
num_support = num_supports[i]
for node in nodes:
one_neighbour = neighbour_list[node]
if len(one_neighbour) >= num_support:
neighbours.append(np.random.choice(neighbour_list[node], num_support, replace=False).tolist())
else:
neighbours.append(np.random.choice(neighbour_list[node], num_support, replace=True).tolist())
layer_neighbours[k - i] = neighbours
nodes = sum(neighbours, [])
return layer_neighbours
def get_nodes_features(nodes, features_embedding, std=True):
embedding = features_embedding[nodes]
if std:
embedding = embedding / embedding.sum(axis=1)
return embedding
if __name__ == '__main__':
train_nodes, train_y, val_nodes, val_y, test_nodes, test_y, graph, features = data_split()
pickle.dump((train_nodes, train_y), open(os.path.join(BASIC_PATH, "./data/train.pkl"), "wb"))
pickle.dump((val_nodes, val_y), open(os.path.join(BASIC_PATH, "./data/val.pkl"), "wb"))
pickle.dump((test_nodes, test_y), open(os.path.join(BASIC_PATH, "./data/test.pkl"), "wb"))
pickle.dump(graph, open(os.path.join(BASIC_PATH, "./data/graph.pkl"), "wb"))
pickle.dump(features, open(os.path.join(BASIC_PATH, "./data/features.pkl"), "wb"))
以上代码将数据集的前1000作为训练,后1000作为测试,中间作为验证,保证节点索引顺序和节点特征向量顺序一致,然后定义了sample
和get_nodes_features
函数,sample目的是输入一组节点,得到他的1跳,2跳节点,get_nodes_features目的是将节点索引转化为节点特征矩阵。
模型部分
模型部分整体参考了https://github.com/williamleif/GraphSAGE,对其进行了简化,具体表现为
-
固定卷积层=2
:大部分情况下卷积层就是2,只要这个一确定,模型层的代码相当好写 -
将采样部分从模型层解耦出来
:源代码中作者将采样部分写入模型层在tensor中计算,好处是这样只需要传递节点索引,可以设置更灵活的卷积层在模型层做采样处理,坏处是他把全局的邻居列表传入了模型做初始化,因此在预测阶段新的邻居列表无法传入。 -
固定聚合方式为MEAN
:先以MEAN跑通流程,其他聚合模式以后再研究。
同时增加了一些新的要素,具体如下:
-
早停
:使用验证集早停,方式过拟合。 -
学习率衰减
:学习率随着训练步长衰减。 -
检查点和pb文件保存
:增加工程部分,能真正实现模型部署。
代码如下,采用标准的tensorflow 1.x风格
import numpy as np
import tensorflow as tf
def glorot(shape, name=None):
init_range = np.sqrt(6.0 / (shape[0] + shape[1]))
initial = tf.random_uniform(shape, minval=-init_range, maxval=init_range, dtype=tf.float32)
return tf.Variable(initial, name=name)
def zeros(shape, name=None):
initial = tf.zeros(shape, dtype=tf.float32)
return tf.Variable(initial, name=name)
class GraphSageGCN(object):
def __init__(self, num_class, feature_size, num_supports_1=10, num_supports_2=10, dim_1=128, dim_2=128,
learning_rate=0.01,
weight_decay=0.01, decay_learning_rate=0.9, concat=True):
# 中心节点
self.input_self = tf.placeholder(tf.float32, [None, feature_size], name="input_self")
# 1跳
self.input_neigh_1 = tf.placeholder(tf.float32, [None, num_supports_1, feature_size], name="input_neigh_1")
# 2跳
self.input_neigh_2 = tf.placeholder(tf.float32, [None, num_supports_1, num_supports_2, feature_size], name="input_neigh_2")
self.input_y = tf.placeholder(tf.int64, [None, num_class])
self.dropout_keep_prob = tf.placeholder(tf.float32, name="dropout_keep_prob")
self.global_step = tf.Variable(0, name="global_step", trainable=False)
# 第1阶卷积
with tf.name_scope('convolution_1'):
# 权重共享
self_weight_1 = glorot([feature_size, dim_1], name="self_weight")
neigh_weight_1 = glorot([feature_size, dim_1], name="neigh_weight")
# 1层的1跳聚合
neigh_vec_1_1 = tf.nn.dropout(self.input_neigh_1, self.dropout_keep_prob)
self_vec_1_1 = tf.nn.dropout(self.input_self, self.dropout_keep_prob)
neigh_means_1_1 = tf.reduce_mean(neigh_vec_1_1, axis=1)
from_neighs_1_1 = tf.matmul(neigh_means_1_1, neigh_weight_1)
from_self_1_1 = tf.matmul(self_vec_1_1, self_weight_1)
output_1_1 = tf.nn.relu(tf.concat([from_self_1_1, from_neighs_1_1], axis=1))
# 1层的2跳聚合
neigh_vec_1_2 = tf.nn.dropout(self.input_neigh_2, self.dropout_keep_prob)
self_vec_1_2 = tf.nn.dropout(self.input_neigh_1, self.dropout_keep_prob)
# [None, feature_size] => [None, None, feature_size]
neigh_means_1_2 = tf.reduce_mean(tf.reshape(neigh_vec_1_2, [-1, num_supports_2, feature_size]), axis=1)
from_neighs_1_2 = tf.matmul(neigh_means_1_2, neigh_weight_1)
from_self_1_2 = tf.matmul(tf.reshape(self_vec_1_2, [-1, feature_size]), self_weight_1)
output_1_2 = tf.nn.relu(tf.concat([from_self_1_2, from_neighs_1_2], axis=1))
# 第二层卷积
with tf.name_scope("convolution_2"):
self_weight_2 = glorot([dim_1 * 2, dim_2], name="self_weight")
neigh_weight_2 = glorot([dim_1 * 2, dim_2], name="neigh_weight")
neigh_vec_2_1 = tf.nn.dropout(output_1_2, self.dropout_keep_prob)
self_vec_2_1 = tf.nn.dropout(output_1_1, self.dropout_keep_prob)
# [None * num_supports[0], dim_1 * 2] => [None, num_supports[0], dim_1 * 2]
neigh_means_2_1 = tf.reduce_mean(tf.reshape(neigh_vec_2_1, [-1, num_supports_1, dim_1 * 2]), axis=1)
from_neighs_2_1 = tf.matmul(neigh_means_2_1, neigh_weight_2)
from_self_2_1 = tf.matmul(self_vec_2_1, self_weight_2)
output_2_1 = tf.concat([from_self_2_1, from_neighs_2_1], axis=1)
# dense
with tf.name_scope("dense"):
dense_weight = tf.get_variable('weights', shape=(dim_2 * 2, num_class),
dtype=tf.float32,
initializer=tf.contrib.layers.xavier_initializer())
dense_bias = zeros([num_class], name='bias')
dense_input = tf.nn.dropout(output_2_1, self.dropout_keep_prob)
output = tf.matmul(dense_input, dense_weight) + dense_bias
# softmax
with tf.name_scope("softmax"):
self.probs = tf.nn.softmax(output, dim=1, name="probs")
self.accuracy = tf.reduce_mean(
tf.cast(tf.equal(tf.arg_max(self.probs, 1), tf.arg_max(self.input_y, 1)), dtype=tf.float32))
# loss
with tf.name_scope('loss'):
# 交叉熵
self.loss = tf.reduce_mean(tf.nn.softmax_cross_entropy_with_logits_v2(logits=output, labels=self.input_y))
# 正则
self.loss += weight_decay * (tf.nn.l2_loss(self_weight_1) + tf.nn.l2_loss(neigh_weight_1) + tf.nn.l2_loss(
self_weight_2) + tf.nn.l2_loss(neigh_weight_2) + tf.nn.l2_loss(dense_weight))
# optimizer
with tf.name_scope("optimizer"):
if decay_learning_rate:
learning_rate = tf.train.exponential_decay(learning_rate, self.global_step, 100, decay_learning_rate)
optimizer = tf.train.AdamOptimizer(learning_rate=learning_rate)
self.train_step = optimizer.minimize(self.loss, global_step=self.global_step)
with tf.name_scope("summaries"):
tf.summary.scalar("loss", self.loss)
tf.summary.scalar("accuracy", self.accuracy)
self.summary_op = tf.summary.merge_all()
可见模型的训练参数就是两层卷积的W和全连接的W,其中设置了3个输入(dropout除外),分别是中心节点特征矩阵
,一跳节点特征矩阵
,2跳节点特征矩阵
,维度分别是2
,3
,4
,全部在model层中写死。保证这三个输入的第一维度值大小是一致的都是batch_size,否则在tensorflow_model_server中无法部署。
训练部分
训练部分主要包括两段内容:
-
get_batch
:批次函数,实现GraphSAGE的小批量训练,采用的按照批次滑动采样,再重复样本,最终返回的是迭代器,内容包括中心节点索引,1条索引,2条索引。 -
train_main
:训练函数,包括模型初始化,将采样节点转化为特征向量,记录训练过程和验证过程,早停的实现。其中早停是连续5次没有出现新的loss低点就停止,每次验证集loss下降就保存一次检查点,最终取最新的检查点,检查点只保存1份。
import sys
import os
import pickle
import shutil
import random
import time
sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
DATA_PATH = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
ROOT_PATH = os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
import numpy as np
import tensorflow as tf
from tensorflow.python.saved_model import tag_constants
from model import GraphSageGCN
from utils.config import get_string
from preprocessing import sample, get_nodes_features
(train_nodes, train_y) = pickle.load(open(os.path.join(DATA_PATH, get_string("train_data_path")), "rb"))
(val_nodes, val_y) = pickle.load(open(os.path.join(DATA_PATH, get_string("val_data_path")), "rb"))
(test_nodes, test_y) = pickle.load(open(os.path.join(DATA_PATH, get_string("test_data_path")), "rb"))
neighbour_list = pickle.load(open(os.path.join(DATA_PATH, get_string("neighbour_data_path")), "rb"))
nodes_features = pickle.load(open(os.path.join(DATA_PATH, get_string("feature_data_path")), "rb"))
features_size = nodes_features.shape[1]
def get_batch(epoches, batch_size, nodes, labels, neighbours, features, layer1_supports=10, layer2_supports=5):
for epoch in range(epoches):
tmp = list(zip(nodes, labels))
random.shuffle(tmp)
nodes, labels = zip(*tmp)
for batch in range(0, len(nodes), batch_size):
if batch + batch_size < len(nodes):
batch_nodes = nodes[batch: (batch + batch_size)]
batch_labels = labels[batch: (batch + batch_size)]
else:
batch_nodes = nodes[batch: len(nodes)]
batch_labels = labels[batch: len(nodes)]
# 得到训练集的1跳2跳
layer_neighbours = sample(batch_nodes, neighbours, num_supports=[layer2_supports, layer1_supports])
# 所有节点的embedding
input_x = get_nodes_features(list(batch_nodes), features)
input_x_1 = get_nodes_features(sum(layer_neighbours[2], []), features)
input_x_2 = get_nodes_features(sum(layer_neighbours[1], []), features)
yield [epoch, input_x, input_x_1, input_x_2, batch_labels]
def train_main():
tf.reset_default_graph()
model = GraphSageGCN(num_class=7, feature_size=1433,
num_supports_1=int(get_string("layer2_supports")),
num_supports_2=int(get_string("layer1_supports")),
decay_learning_rate=float(get_string("decay_learning_rate")),
learning_rate=float(get_string("learning_rate")),
weight_decay=float(get_string("weight_decay")))
saver = tf.train.Saver(tf.global_variables(), max_to_keep=1)
with tf.Session() as sess:
init_op = tf.group(tf.global_variables_initializer())
sess.run(init_op)
shutil.rmtree(os.path.join(ROOT_PATH, "./summary"), ignore_errors=True)
writer = tf.summary.FileWriter(os.path.join(ROOT_PATH, "./summary"), sess.graph)
batches = get_batch(int(get_string("epoches")), int(get_string("batch_size")), train_nodes, train_y,
neighbour_list, nodes_features, layer1_supports=int(get_string("layer1_supports")),
layer2_supports=int(get_string("layer2_supports")))
# 验证数据
layer_neighbours = sample(val_nodes, neighbour_list,
num_supports=[int(get_string("layer2_supports")), int(get_string("layer1_supports"))])
val_input_x = get_nodes_features(val_nodes, nodes_features)
val_input_x_1 = get_nodes_features(sum(layer_neighbours[2], []), nodes_features)
val_input_x_2 = get_nodes_features(sum(layer_neighbours[1], []), nodes_features)
val_feed_dict = {model.input_self: val_input_x,
model.input_neigh_1: val_input_x_1.A.reshape(-1, int(get_string("layer1_supports")),
features_size),
model.input_neigh_2: val_input_x_2.A.reshape(-1, int(get_string("layer1_supports")),
int(get_string("layer2_supports")),
features_size),
model.input_y: val_y,
model.dropout_keep_prob: 1}
val_loss_list = []
for batch in batches:
epoch, input_x, input_x_1, input_x_2, input_y = batch
feed_dict = {model.input_self: input_x,
model.input_neigh_1: input_x_1.A.reshape(-1, int(get_string("layer1_supports")),
features_size),
model.input_neigh_2: input_x_2.A.reshape(-1, int(get_string("layer1_supports")),
int(get_string("layer2_supports")), features_size),
model.input_y: input_y,
model.dropout_keep_prob: float(get_string("dropout_keep_prob"))}
_, step, loss_val, acc_val, merged = sess.run(
[model.train_step, model.global_step, model.loss, model.accuracy, model.summary_op],
feed_dict=feed_dict)
writer.add_summary(merged, step)
if step % 5 == 0:
print("epoch:", epoch + 1, "step:", step, "loss:", loss_val, "accuracy:", acc_val)
if step % 20 == 0:
loss_val, acc_val = sess.run([model.loss, model.accuracy], feed_dict=val_feed_dict)
print("{:-^30}".format("evaluation"))
print("[evaluation]", "loss:", loss_val, "accuracy:", acc_val)
# 计算当前loss相比之前的最有loss下降多少
diff = (loss_val - min(val_loss_list)) if len(val_loss_list) else 0
val_loss_list.append(loss_val)
print("本轮loss比之前最小loss{}:{}, 当前最小loss: {}"
.format("上升" if diff > 0 else "下降", abs(diff), min(val_loss_list)))
if diff < 0:
saver.save(sess, os.path.join(ROOT_PATH, get_string("checkpoint_path")))
print("[save checkpoint]")
print("-" * 40)
if early_stop(val_loss_list, windows=int(get_string("early_stop_windows"))):
print("{:-^30}".format("early stop!"))
break
def early_stop(loss_list, windows=5):
if len(loss_list) <= windows:
return False
latest_loss = loss_list[-windows:]
previous_loss = loss_list[:-windows]
min_previous_loss = min(previous_loss)
min_latest_loss = min(latest_loss)
if min_latest_loss > min_previous_loss:
return True
return False
运行之后打印的输出如下
epoch: 12 step: 185 loss: 0.5722251 accuracy: 0.953125
epoch: 12 step: 190 loss: 0.52046406 accuracy: 0.984375
epoch: 13 step: 195 loss: 0.57312083 accuracy: 0.9375
epoch: 13 step: 200 loss: 0.482001 accuracy: 1.0
----------evaluation----------
[evaluation] loss: 0.85649866 accuracy: 0.8629944
本轮loss比之前最小loss上升:0.038332998752593994, 当前最小loss: 0.81816565990448
----------------------------------------
epoch: 13 step: 205 loss: 0.56127137 accuracy: 0.953125
epoch: 14 step: 210 loss: 0.5034015 accuracy: 0.984375
epoch: 14 step: 215 loss: 0.4905039 accuracy: 0.984375
epoch: 14 step: 220 loss: 0.5316473 accuracy: 0.96875
----------evaluation----------
[evaluation] loss: 0.8504291 accuracy: 0.8615819
本轮loss比之前最小loss上升:0.03226345777511597, 当前最小loss: 0.81816565990448
----------------------------------------
---------early stop!----------
最终训练集准确率能达到1,验证集准确率有0.86。
测试部分
测试部分读取最新的检查点,使用get_tensor_by_name
拿到输入tensor和输出tensor,用新的test数据灌入输入tensor拿到预测tensor。
def test_main():
layer_neighbours = sample(test_nodes, neighbour_list,
num_supports=[int(get_string("layer2_supports")), int(get_string("layer1_supports"))])
test_input_x = get_nodes_features(test_nodes, nodes_features)
test_input_x_1 = get_nodes_features(sum(layer_neighbours[2], []), nodes_features)
test_input_x_2 = get_nodes_features(sum(layer_neighbours[1], []), nodes_features)
tf.reset_default_graph()
with tf.Session() as sess:
last_ckpt = tf.train.latest_checkpoint(
os.path.join(ROOT_PATH, "/".join(get_string("checkpoint_path").split("/")[:-1])))
print("读取ckpt: {}".format(last_ckpt))
saver = tf.train.import_meta_graph("{}.meta".format(last_ckpt))
saver.restore(sess, last_ckpt)
graph = tf.get_default_graph()
# get tensor
input_self = graph.get_tensor_by_name("input_self:0")
input_neigh_1 = graph.get_tensor_by_name("input_neigh_1:0")
input_neigh_2 = graph.get_tensor_by_name("input_neigh_2:0")
dropout_keep_prob = graph.get_tensor_by_name("dropout_keep_prob:0")
pred = graph.get_tensor_by_name("softmax/probs:0")
prediction = sess.run(pred, feed_dict={input_self: test_input_x,
input_neigh_1: test_input_x_1.A.reshape(-1, int(get_string("layer1_supports")), features_size),
input_neigh_2: test_input_x_2.A.reshape(-1, int(get_string("layer1_supports")), int(get_string("layer2_supports")), features_size),
dropout_keep_prob: 1.0})
hit = np.equal(np.argmax(prediction, axis=1), np.argmax(test_y, axis=1))
accuracy = hit.sum() / len(hit)
print("[test]:", accuracy)
执行打印的输入如下
读取ckpt: /home/myproject/GRAPHSAGE_CORA/./ckpt/ckpt
[test]: 0.834
测试集准备率有0.834
tensorboard查看训练标量
tensorboard简单查看以下训练过程,用以做ppt使用,切到summary目录下
tensorboard --logdir `pwd`
打开127.0.0.1:6006
accuracy持续上升,loss持续下降
模型保存为pb
模型保存为pb也是拿的最新的检查点文件,都是tensorflow的标准代码
def save_pb():
# 模型保存
pb_num = str(int(time.time()))
pb_path = os.path.join(ROOT_PATH, get_string("pb_path"), pb_num)
shutil.rmtree(pb_path, ignore_errors=True)
tf.reset_default_graph()
with tf.Session() as sess:
last_ckpt = tf.train.latest_checkpoint(
os.path.join(ROOT_PATH, "/".join(get_string("checkpoint_path").split("/")[:-1])))
print("读取ckpt: {}".format(last_ckpt))
saver = tf.train.import_meta_graph("{}.meta".format(last_ckpt))
saver.restore(sess, last_ckpt)
graph = tf.get_default_graph()
# get tensor
input_self = graph.get_tensor_by_name("input_self:0")
input_neigh_1 = graph.get_tensor_by_name("input_neigh_1:0")
input_neigh_2 = graph.get_tensor_by_name("input_neigh_2:0")
dropout_keep_prob = graph.get_tensor_by_name("dropout_keep_prob:0")
pred = graph.get_tensor_by_name("softmax/probs:0")
builder = tf.saved_model.builder.SavedModelBuilder(pb_path)
inputs = {'input_self': tf.saved_model.utils.build_tensor_info(input_self),
'input_neigh_1': tf.saved_model.utils.build_tensor_info(input_neigh_1),
'input_neigh_2': tf.saved_model.utils.build_tensor_info(input_neigh_2),
'dropout_keep_prob': tf.saved_model.utils.build_tensor_info(dropout_keep_prob),
}
outputs = {'output': tf.saved_model.utils.build_tensor_info(pred)}
signature = tf.saved_model.signature_def_utils.build_signature_def(
inputs=inputs,
outputs=outputs,
method_name=tf.saved_model.signature_constants.PREDICT_METHOD_NAME)
builder.add_meta_graph_and_variables(sess, [tag_constants.SERVING], {'my_signature': signature})
builder.save()
print("pb文件保存完成:", pb_num)
执行完毕后tfserving目录下会有多个时间戳的模型文件
root@ubuntu:~/myproject/GRAPHSAGE_CORA/tfserving# tree
.
└── 1641799963
├── saved_model.pb
└── variables
├── variables.data-00000-of-00001
└── variables.index
模型server API服务
先使用saved_model_cli
起一个测试服务,先查看一下pb的参数信息
root@ubuntu:~/myproject/GRAPHSAGE_CORA/tfserving# saved_model_cli show --all --dir 1641799963
输出如下,可以看到4个input信息以及维度,一个output信息
signature_def['my_signature']:
The given SavedModel SignatureDef contains the following input(s):
inputs['dropout_keep_prob'] tensor_info:
dtype: DT_FLOAT
shape: unknown_rank
name: dropout_keep_prob:0
inputs['input_neigh_1'] tensor_info:
dtype: DT_FLOAT
shape: (-1, 10, 1433)
name: input_neigh_1:0
inputs['input_neigh_2'] tensor_info:
dtype: DT_FLOAT
shape: (-1, 10, 10, 1433)
name: input_neigh_2:0
inputs['input_self'] tensor_info:
dtype: DT_FLOAT
shape: (-1, 1433)
name: input_self:0
The given SavedModel SignatureDef contains the following output(s):
outputs['output'] tensor_info:
dtype: DT_FLOAT
shape: (-1, 7)
name: softmax/probs:0
Method name is: tensorflow/serving/predict
可以看到需要输入的参数和维度,类型要求,下一步启动docker的serving服务
docker run --rm \
-p 13713:8501 \
-v /home/myproject/GRAPHSAGE_CORA/tfserving/:/models/graphsage_cora/ \
-e MODEL_NAME=graphsage_cora \
--name graphsage_api \
tensorflow/serving
注意MODEL_NAME
和挂载到容器内的models下的目录名称保持一致,写一个Python脚本测试一下api调用
import os
import pickle
import sys
sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
import requests
from preprocessing import sample, get_nodes_features
from utils.config import get_string
DATA_PATH = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
neighbour_list = pickle.load(open(os.path.join(DATA_PATH, get_string("neighbour_data_path")), "rb"))
nodes_features = pickle.load(open(os.path.join(DATA_PATH, get_string("feature_data_path")), "rb"))
if __name__ == '__main__':
nodes = [2555]
layer_neighbours = sample(nodes, neighbour_list,
num_supports=[int(get_string("layer2_supports")), int(get_string("layer1_supports"))])
test_input_x = get_nodes_features(nodes, nodes_features)
test_input_x_1 = get_nodes_features(sum(layer_neighbours[2], []), nodes_features)
test_input_x_2 = get_nodes_features(sum(layer_neighbours[1], []), nodes_features)
res = requests.post("http://127.0.0.1:13713/v1/models/graphsage_cora:predict", json={"instances": [{
"input_self": test_input_x.A[0].tolist(),
"input_neigh_1": test_input_x_1.A.reshape(-1, 10, 1433)[0].tolist(),
"input_neigh_2": test_input_x_2.A.reshape(-1, 10, 10, 1433)[0].tolist(),
"dropout_keep_prob": 1.0
}], "signature_name": "my_signature"})
print(res.json())
输出如下
{'predictions': [[0.0151401442, 0.00185001828, 0.000755791727, 0.97733295, 0.00436462974, 0.000487773854, 6.87406064e-05]]}
结果是第2555个节点预测是第4分类。