从协议缓冲区 / tfrecords 解码字节时的动态 dtype

Question

我正在尝试在 tensorflow 的 python 接口中读取 TFRecords 文件。文件中的每个示例都包含一个 n 维张量及其原始数据类型。 n 维张量在保存之前被序列化为字节。在读取 TFRecords 文件时，我想根据每个张量的数据类型对其进行解码。但是，当我尝试这个时，我遇到了错误，因为 out_type 的 tf.io.decode_raw 不期望张量。我在下面提供了一个示例。如何根据示例中存储的 dtype 动态分配 out_type？

import numpy as np
import tensorflow as tf

def _bytes_feature(value):
    """Returns a bytes_list from a string / byte."""
    if isinstance(value, type(tf.constant(0))):
        value = value.numpy() # BytesList won't unpack a string from an EagerTensor.
    return tf.train.Feature(bytes_list=tf.train.BytesList(value=[value]))

def serialize(x):
    feature = {
        "value": _bytes_feature(x.tobytes()),
        "dtype": _bytes_feature(x.dtype.name.encode('utf-8'))}

    example = tf.train.Example(features=tf.train.Features(feature=feature))
    return example.SerializeToString()

def parse(serialized):
    features = {
        "value": tf.io.FixedLenFeature(shape=[], dtype=tf.string),
        "dtype": tf.io.FixedLenFeature(shape=[], dtype=tf.string)}
    return tf.io.parse_single_example(serialized, features=features)

x = np.random.random_sample((10, 10, 10)).astype(np.float32)
serialized = serialize(x)
parsed = parse(serialized)

# This line causes the error.
tf.io.decode_raw(parsed["value"], out_type=parsed["dtype"])

# This works.
tf.io.decode_raw(parsed["value"], out_type="float32")

Answer 1

tf.io.decode_raw(parsed["value"], out_type=parsed["dtype"].numpy().decode())

从原始的 out_type 张量获取它的值 numpy=b"float32"。然后将字节解码为字符串，这将给出“float32”（或其他动态 dtype）。