python - 保存模型 DNNEstimator Tensorflow

Question

我有这个训练算法，但我无法保存（导出 SavedModel），我查看了 tensorflow 文档但我不明白。

如何保存在函数返回时返回的训练模型？

import os
import numpy as np
import pandas as pd
import tensorflow as tf
import tensorflow_hub as hub
from sklearn.preprocessing import MultiLabelBinarizer

def classifica_role(vaga, texto, role):

    train_size = int(len(vaga[texto]) * .8)

    train_texto = vaga[texto][:train_size].astype('str')
    train_role = vaga[role][:train_size]

    test_texto = vaga[texto][train_size:].astype('str')
    test_role = vaga[role][train_size:]

    encoder = MultiLabelBinarizer()
    encoder.fit_transform(train_role)
    train_encoded = encoder.transform(train_role)
    test_encoded = encoder.transform(test_role)
    num_classes = len(encoder.classes_)

    description_embeddings = hub.text_embedding_column("descricao", module_spec="https://tfhub.dev/google/universal-sentence-encoder/2", trainable=False)

    multi_label_head = tf.contrib.estimator.multi_label_head(num_classes, loss_reduction=tf.losses.Reduction.SUM_OVER_BATCH_SIZE)

    features = {"descricao": np.array(train_texto).astype(np.str)}
    labels = np.array(train_encoded).astype(np.int32)
    train_input_fn = tf.estimator.inputs.numpy_input_fn(features, labels, shuffle=True, batch_size=32, num_epochs=25)
    estimator = tf.estimator.DNNEstimator(head = multi_label_head, hidden_units = [64,10], feature_columns = [description_embeddings])

    return estimator.train(input_fn=train_input_fn)

标签： pythonpython-3.xtensorflowmachine-learninganaconda