TensorFlow tf.map_fn 移除一个维度

Question

我正在尝试使用 defaultdict 重新映射输入张量中的值。

class MyDataSet(object):
    def __init__(self):
        self.class_map = MyDataSet.remap_class()

    @staticmethod
    def remap_class():
        class_remap = defaultdict(lambda: 11)
        class_remap[128] = 0  
        class_remap[130] = 1  
        class_remap[132] = 2
        # ...

    def parser(self, serialized_example):
        features = tf.parse_single_example(
            serialized_example,
            features={
                'image': tf.FixedLenFeature([], tf.string),
                'label': tf.FixedLenFeature([], tf.string),
            })
        label = tf.decode_raw(features['label'], tf.uint8)
        label.set_shape([256 * 512])
        label = tf.cast(tf.reshape(label, [256, 512]), tf.int32)

        output_label = tf.map_fn(lambda x: self.class_map(x), label)

    #...
    dataset = tf.data.TFRecordDataset(filenames).repeat()
    dataset = dataset.map(self.parser, num_parallel_calls=batch_size)

标签形状是 (256,512) 但 output_label 形状是 (256,)。如果我尝试使用

更改 output_label

output_label = tf.reshape(output_label, [256, 512])

我得到了异常

ValueError: Cannot reshape a tensor with 256 elements to shape [256,512] (131072 elements) for 'Reshape_2' (op: 'Reshape') with input shapes: [256], [2] and with input tensors computed as partial shapes: input[1] = [256,512].

如果我尝试用

更改 output_label

output_label.set_shape([256, 512])

我得到了异常

ValueError: Shapes (256,) and (256, 512) must have the same rank

如何在 output_label 中映射值并保持与 label 中相同的形状？

Answer 1

解决该问题的方法是对展平张量进行操作。所以改变：

    def parser(self, serialized_example):
        features = tf.parse_single_example(
            serialized_example,
            features={
                'image': tf.FixedLenFeature([], tf.string),
                'label': tf.FixedLenFeature([], tf.string),
            })
        label = tf.decode_raw(features['label'], tf.uint8)
        label.set_shape([256 * 512])
        label = tf.cast(tf.reshape(label, [256, 512]), tf.int32)

        output_label = tf.map_fn(lambda x: self.class_map(x), label)

到：

    def parser(self, serialized_example):
        features = tf.parse_single_example(
            serialized_example,
            features={
                'image': tf.FixedLenFeature([], tf.string),
                'label': tf.FixedLenFeature([], tf.string),
            })
        label = tf.decode_raw(features['label'], tf.uint8)
        label.set_shape([256 * 512])

        output_label = tf.map_fn(lambda x: self.class_map(x), label)
        label = tf.cast(tf.reshape(label, [256, 512]), tf.int32)
        output_label = tf.cast(tf.reshape(output_label, [256, 512]), tf.int32)