更新时间:2023-12-02 18:41:58
I found a solution here.
def _make_input_parser(with_target=True):
"""Returns a parser func according to file_type, task_type and target.
Need to set record_default for last column to integer instead of float in
case of classification tasks.
Args:
with_target (boolean): Pass label or not.
Returns:
It returns a parser.
"""
def _decode_csv(line):
"""Takes the string input tensor and parses it to feature dict and target.
All the columns except the first one are treated as feature column. The
first column is expected to be the target.
Only returns target for if with_target is True.
Args:
line: csv rows in tensor format.
Returns:
features: A dictionary of features with key as "column_names" from
self._column_header.
target: tensor of target values which is the first column of the file.
This will only be returned if with_target==True.
"""
column_header = column_names if with_target else column_names[:4]
record_defaults = [[0.] for _ in xrange(len(column_names) - 1)]
# Pass label as integer.
if with_target:
record_defaults.append([0])
columns = tf.decode_csv(line, record_defaults=record_defaults)
features = dict(zip(column_header, columns))
target = features.pop(column_names[4]) if with_target else None
return features, target
return _decode_csv
def serving_input_receiver_fn():
"""This is used to define inputs to serve the model.
Returns:
A ServingInputReciever object.
"""
csv_row = tf.placeholder(shape=[None], dtype=tf.string)
features, _ = _make_input_parser(with_target=False)(csv_row)
return tf.estimator.export.ServingInputReceiver(features,
{'csv_row': csv_row})